tasks package

Subpackages

• tasks.DivaHisDB package
  • Submodules
  • tasks.DivaHisDB.semantic_segmentation_cropped module
  • Module contents
• tasks.RGB package
  • Submodules
  • tasks.RGB.semantic_segmentation module
  • tasks.RGB.semantic_segmentation_cropped module
  • Module contents
• tasks.classification package
  • Submodules
  • tasks.classification.classification module
  • Module contents
• tasks.utils package
  • Submodules
  • tasks.utils.outputs module
  • tasks.utils.task_utils module
  • Module contents

Submodules

tasks.base_task module

class AbstractTask(model: ~typing.Optional[~torch.nn.modules.module.Module] = None, loss_fn: ~typing.Optional[~typing.Union[~typing.Callable, ~typing.Mapping, ~typing.Sequence]] = None, optimizer: ~typing.Union[~typing.Type[~torch.optim.optimizer.Optimizer], ~torch.optim.optimizer.Optimizer] = <class 'torch.optim.adam.Adam'>, optimizer_kwargs: ~typing.Optional[~typing.Dict[str, ~typing.Any]] = None, scheduler: ~typing.Optional[~typing.Union[~typing.Type[~torch.optim.lr_scheduler._LRScheduler], str, ~torch.optim.lr_scheduler._LRScheduler]] = None, scheduler_kwargs: ~typing.Optional[~typing.Dict[str, ~typing.Any]] = None, metric_train: ~typing.Optional[~torchmetrics.collections.MetricCollection] = None, metric_val: ~typing.Optional[~torchmetrics.collections.MetricCollection] = None, metric_test: ~typing.Optional[~torchmetrics.collections.MetricCollection] = None, confusion_matrix_val: ~typing.Optional[bool] = False, confusion_matrix_test: ~typing.Optional[bool] = False, confusion_matrix_log_every_n_epoch: ~typing.Optional[int] = 1, lr: float = 0.001, test_output_path: ~typing.Optional[~typing.Union[str, ~pathlib.Path]] = 'test_output', predict_output_path: ~typing.Optional[~typing.Union[str, ~pathlib.Path]] = 'predict_output')

Bases: LightningModule

Inspired by Pytorch Lightning.

A general abstract Task. It provieds the basic functionality for training, validation and testing. A step method is provided which can be overwritten for custom behavior. The step method is called in the training, validation and testing loop. The step method should return a dictionary with the following keys:

• OutputKeys.PREDICTION: The prediction of the model.

• OutputKeys.LOSS: The loss of the model.

• OutputKeys.LOG: A dictionary with all the logs. The keys are the metric names and the values are the

  metric values.

• OutputKeys.TARGET: The target of the model.

Parameters:

• model (nn.Module) – Composed model to use for the task.

• loss_fn (Union[Callable, Mapping, Sequence]) – Loss function for training

• optimizer (Union[Type[torch.optim.Optimizer], torch.optim.Optimizer]) – Optimizer to use for training, defaults to torch.optim.Adam.

• optimizer_kwargs (Optional[Dict[str, Any]]) – Keyword arguments to pass to the optimizer.

• scheduler (Optional[Union[Type[_LRScheduler], str, _LRScheduler]]) – Learning rate scheduler to use for training, defaults to None.

• scheduler_kwargs (Optional[Dict[str, Any]]) – Keyword arguments to pass to the scheduler.

• metric_train (Optional[MetricCollection]) – Metrics to compute for training.

• metric_val (Optional[MetricCollection]) – Metrics to compute for evaluation.

• metric_test (Optional[MetricCollection]) – Metrics to compute for testing.

• confusion_matrix_val (bool) – Whether to compute the confusion matrix for the validation set.

• confusion_matrix_test (bool) – Whether to compute the confusion matrix for the test set.

• confusion_matrix_log_every_n_epoch (int) – How often to compute the confusion matrix.

• lr (float) – Learning rate to use for training, defaults to 5e-5.

• test_output_path (Union[str, Path]) – Path relative to the normal output folder where to save the test output

• predict_output_path (Union[str, Path]) – Path relative to the normal output folder where to save the predict output

configure_optimizers() → Union[Optimizer, Tuple[List[Optimizer], List[_LRScheduler]]]

Choose what optimizers and learning-rate schedulers to use in your optimization. Normally you’d need one. But in the case of GANs or similar you might have multiple.

Returns:

Any of these 6 options.

• Single optimizer.

• List or Tuple of optimizers.

• Two lists - The first list has multiple optimizers, and the second has multiple LR schedulers (or multiple lr_scheduler_config).

• Dictionary, with an "optimizer" key, and (optionally) a "lr_scheduler" key whose value is a single LR scheduler or lr_scheduler_config.

• Tuple of dictionaries as described above, with an optional "frequency" key.

• None - Fit will run without any optimizer.

The lr_scheduler_config is a dictionary which contains the scheduler and its associated configuration. The default configuration is shown below.

lr_scheduler_config = {
    # REQUIRED: The scheduler instance
    "scheduler": lr_scheduler,
    # The unit of the scheduler's step size, could also be 'step'.
    # 'epoch' updates the scheduler on epoch end whereas 'step'
    # updates it after a optimizer update.
    "interval": "epoch",
    # How many epochs/steps should pass between calls to
    # `scheduler.step()`. 1 corresponds to updating the learning
    # rate after every epoch/step.
    "frequency": 1,
    # Metric to to monitor for schedulers like `ReduceLROnPlateau`
    "monitor": "val_loss",
    # If set to `True`, will enforce that the value specified 'monitor'
    # is available when the scheduler is updated, thus stopping
    # training if not found. If set to `False`, it will only produce a warning
    "strict": True,
    # If using the `LearningRateMonitor` callback to monitor the
    # learning rate progress, this keyword can be used to specify
    # a custom logged name
    "name": None,
}

When there are schedulers in which the .step() method is conditioned on a value, such as the torch.optim.lr_scheduler.ReduceLROnPlateau scheduler, Lightning requires that the lr_scheduler_config contains the keyword "monitor" set to the metric name that the scheduler should be conditioned on.

Metrics can be made available to monitor by simply logging it using self.log('metric_to_track', metric_val) in your LightningModule.

Note

The frequency value specified in a dict along with the optimizer key is an int corresponding to the number of sequential batches optimized with the specific optimizer. It should be given to none or to all of the optimizers. There is a difference between passing multiple optimizers in a list, and passing multiple optimizers in dictionaries with a frequency of 1:

• In the former case, all optimizers will operate on the given batch in each optimization step.

• In the latter, only one optimizer will operate on the given batch at every step.

This is different from the frequency value specified in the lr_scheduler_config mentioned above.

def configure_optimizers(self):
    optimizer_one = torch.optim.SGD(self.model.parameters(), lr=0.01)
    optimizer_two = torch.optim.SGD(self.model.parameters(), lr=0.01)
    return [
        {"optimizer": optimizer_one, "frequency": 5},
        {"optimizer": optimizer_two, "frequency": 10},
    ]

In this example, the first optimizer will be used for the first 5 steps, the second optimizer for the next 10 steps and that cycle will continue. If an LR scheduler is specified for an optimizer using the lr_scheduler key in the above dict, the scheduler will only be updated when its optimizer is being used.

Examples:

# most cases. no learning rate scheduler
def configure_optimizers(self):
    return Adam(self.parameters(), lr=1e-3)

# multiple optimizer case (e.g.: GAN)
def configure_optimizers(self):
    gen_opt = Adam(self.model_gen.parameters(), lr=0.01)
    dis_opt = Adam(self.model_dis.parameters(), lr=0.02)
    return gen_opt, dis_opt

# example with learning rate schedulers
def configure_optimizers(self):
    gen_opt = Adam(self.model_gen.parameters(), lr=0.01)
    dis_opt = Adam(self.model_dis.parameters(), lr=0.02)
    dis_sch = CosineAnnealing(dis_opt, T_max=10)
    return [gen_opt, dis_opt], [dis_sch]

# example with step-based learning rate schedulers
# each optimizer has its own scheduler
def configure_optimizers(self):
    gen_opt = Adam(self.model_gen.parameters(), lr=0.01)
    dis_opt = Adam(self.model_dis.parameters(), lr=0.02)
    gen_sch = {
        'scheduler': ExponentialLR(gen_opt, 0.99),
        'interval': 'step'  # called after each training step
    }
    dis_sch = CosineAnnealing(dis_opt, T_max=10) # called every epoch
    return [gen_opt, dis_opt], [gen_sch, dis_sch]

# example with optimizer frequencies
# see training procedure in `Improved Training of Wasserstein GANs`, Algorithm 1
# https://arxiv.org/abs/1704.00028
def configure_optimizers(self):
    gen_opt = Adam(self.model_gen.parameters(), lr=0.01)
    dis_opt = Adam(self.model_dis.parameters(), lr=0.02)
    n_critic = 5
    return (
        {'optimizer': dis_opt, 'frequency': n_critic},
        {'optimizer': gen_opt, 'frequency': 1}
    )

Note

Some things to know:

• Lightning calls .backward() and .step() on each optimizer as needed.

• If learning rate scheduler is specified in configure_optimizers() with key "interval" (default “epoch”) in the scheduler configuration, Lightning will call the scheduler’s .step() method automatically in case of automatic optimization.

• If you use 16-bit precision (precision=16), Lightning will automatically handle the optimizers.

• If you use multiple optimizers, training_step() will have an additional optimizer_idx parameter.

• If you use torch.optim.LBFGS, Lightning handles the closure function automatically for you.

• If you use multiple optimizers, gradients will be calculated only for the parameters of current optimizer at each training step.

• If you need to control how often those optimizers step or override the default .step() schedule, override the optimizer_step() hook.

forward(x: Any) → Any

Same as torch.nn.Module.forward().

Parameters:

• *args – Whatever you decide to pass into the forward method.

• **kwargs – Keyword arguments are also possible.

Returns:

Your model’s output

predict_step(batch: Any, batch_idx: int, dataloader_idx: Optional[int] = None) → Any

Step function called during predict(). By default, it calls forward(). Override to add any processing logic.

The predict_step() is used to scale inference on multi-devices.

To prevent an OOM error, it is possible to use BasePredictionWriter callback to write the predictions to disk or database after each batch or on epoch end.

The BasePredictionWriter should be used while using a spawn based accelerator. This happens for Trainer(strategy="ddp_spawn") or training on 8 TPU cores with Trainer(accelerator="tpu", devices=8) as predictions won’t be returned.

Example

class MyModel(LightningModule):

    def predict_step(self, batch, batch_idx, dataloader_idx=0):
        return self(batch)

dm = ...
model = MyModel()
trainer = Trainer(accelerator="gpu", devices=2)
predictions = trainer.predict(model, dm)

Parameters:

• batch – Current batch.

• batch_idx – Index of current batch.

• dataloader_idx – Index of the current dataloader.

Returns:

Predicted output

setup(stage: str)

Called at the beginning of fit (train + validate), validate, test, or predict. This is a good hook when you need to build models dynamically or adjust something about them. This hook is called on every process when using DDP.

Parameters:

stage – either 'fit', 'validate', 'test', or 'predict'

Example:

class LitModel(...):
    def __init__(self):
        self.l1 = None

    def prepare_data(self):
        download_data()
        tokenize()

        # don't do this
        self.something = else

    def setup(self, stage):
        data = load_data(...)
        self.l1 = nn.Linear(28, data.num_classes)

step(batch: Any, metric_kwargs: Optional[Dict[str, Dict[str, Any]]] = None) → Union[Dict[OutputKeys, Any], Tuple[Any, Any]]

The training/validation/test step. Override for custom behavior.

Parameters:

• batch (Any) – the batch with the images (x) and the gt (y) in the order (x, y)

• metric_kwargs (Optional[Dict[str, Dict[str, Any]]]) – a dictionary with a entry with the additional arguments (pred and y always provided). e.g. you have two metrics (A, B) and B takes an additional arguments x and y so the dictionary would look like this: {‘B’: {‘x’: ‘value’, ‘y’: ‘value’}}

test_epoch_end(outputs: Any) → None

Called at the end of a test epoch with the output of all test steps.

# the pseudocode for these calls
test_outs = []
for test_batch in test_data:
    out = test_step(test_batch)
    test_outs.append(out)
test_epoch_end(test_outs)

Parameters:

outputs – List of outputs you defined in test_step_end(), or if there are multiple dataloaders, a list containing a list of outputs for each dataloader

Returns:

None

Note

If you didn’t define a test_step(), this won’t be called.

Examples

With a single dataloader:

def test_epoch_end(self, outputs):
    # do something with the outputs of all test batches
    all_test_preds = test_step_outputs.predictions

    some_result = calc_all_results(all_test_preds)
    self.log(some_result)

With multiple dataloaders, outputs will be a list of lists. The outer list contains one entry per dataloader, while the inner list contains the individual outputs of each test step for that dataloader.

def test_epoch_end(self, outputs):
    final_value = 0
    for dataloader_outputs in outputs:
        for test_step_out in dataloader_outputs:
            # do something
            final_value += test_step_out

    self.log("final_metric", final_value)

test_step(batch: Any, batch_idx: int, **kwargs) → None

Operates on a single batch of data from the test set. In this step you’d normally generate examples or calculate anything of interest such as accuracy.

# the pseudocode for these calls
test_outs = []
for test_batch in test_data:
    out = test_step(test_batch)
    test_outs.append(out)
test_epoch_end(test_outs)

Parameters:

• batch – The output of your DataLoader.

• batch_idx – The index of this batch.

• dataloader_id – The index of the dataloader that produced this batch. (only if multiple test dataloaders used).

Returns:

Any of.

• Any object or value

• None - Testing will skip to the next batch

# if you have one test dataloader:
def test_step(self, batch, batch_idx):
    ...


# if you have multiple test dataloaders:
def test_step(self, batch, batch_idx, dataloader_idx=0):
    ...

Examples:

# CASE 1: A single test dataset
def test_step(self, batch, batch_idx):
    x, y = batch

    # implement your own
    out = self(x)
    loss = self.loss(out, y)

    # log 6 example images
    # or generated text... or whatever
    sample_imgs = x[:6]
    grid = torchvision.utils.make_grid(sample_imgs)
    self.logger.experiment.add_image('example_images', grid, 0)

    # calculate acc
    labels_hat = torch.argmax(out, dim=1)
    test_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)

    # log the outputs!
    self.log_dict({'test_loss': loss, 'test_acc': test_acc})

If you pass in multiple test dataloaders, test_step() will have an additional argument. We recommend setting the default value of 0 so that you can quickly switch between single and multiple dataloaders.

# CASE 2: multiple test dataloaders
def test_step(self, batch, batch_idx, dataloader_idx=0):
    # dataloader_idx tells you which dataset this is.
    ...

Note

If you don’t need to test you don’t need to implement this method.

Note

When the test_step() is called, the model has been put in eval mode and PyTorch gradients have been disabled. At the end of the test epoch, the model goes back to training mode and gradients are enabled.

static to_loss_format(x: Tensor, **kwargs) → Tensor

Convert the output of the model to the format needed for the loss function.

Parameters:

• x (torch.Tensor) – the output of the model

• kwargs (Any) – additional arguments

Returns:

the output in the format needed for the loss function

Return type:

torch.Tensor

static to_metrics_format(x: Tensor, **kwargs) → Tensor

Convert the output of the model to the format needed for the metrics.

Parameters:

• x (torch.Tensor) – the output of the model

• kwargs (Any) – additional arguments

Returns:

the output in the format needed for the metrics

Return type:

torch.Tensor

training: bool

training_step(batch: Any, batch_idx: int, **kwargs) → Any

The training step. Calls the step method and logs the metrics and loss.

Parameters:

• batch (Any) – The current batch to train on.

• batch_idx (int) – The index of the current batch.

• kwargs (Any) – Additional arguments.

Returns:

The output of the step method.

Return type:

Any

validation_epoch_end(outputs: Any) → None

Called at the end of the validation epoch with the outputs of all validation steps.

# the pseudocode for these calls
val_outs = []
for val_batch in val_data:
    out = validation_step(val_batch)
    val_outs.append(out)
validation_epoch_end(val_outs)

Parameters:

outputs – List of outputs you defined in validation_step(), or if there are multiple dataloaders, a list containing a list of outputs for each dataloader.

Returns:

None

Note

If you didn’t define a validation_step(), this won’t be called.

Examples

With a single dataloader:

def validation_epoch_end(self, val_step_outputs):
    for out in val_step_outputs:
        ...

With multiple dataloaders, outputs will be a list of lists. The outer list contains one entry per dataloader, while the inner list contains the individual outputs of each validation step for that dataloader.

def validation_epoch_end(self, outputs):
    for dataloader_output_result in outputs:
        dataloader_outs = dataloader_output_result.dataloader_i_outputs

    self.log("final_metric", final_value)

validation_step(batch: Any, batch_idx: int, **kwargs) → None

the validation step. Calls the step method and logs the metrics and loss.

Parameters:

• batch (Any) – The current batch to validate on.

• batch_idx (int) – The index of the current batch.

• kwargs (Any) – Additional arguments.

Returns:

The output of the step method.

Return type:

Any

Module contents