Core (Catalyst abstractions)

There are 3 main abstractions in Catalyst: Experiment, Runner, and Callback.

In general, the Experiment knows what you would like to run, Runner contains all the logic of how to run the experiment, while Callbacks allow you to customize experiment run logic by using specific custom callback functions without changing neither Experiment nor Runner.

Note

To learn more about Catalyst Core concepts, please check out

  • catalyst.core.experiment.IExperiment (docs)

  • catalyst.core.runner.IRunner (docs)

  • catalyst.core.callback.Callback (docs)

Experiment

Experiment in an abstraction that contains information about the experiment – a model, a criterion, an optimizer, a scheduler, and their hyperparameters. It also contains information about the data and transformations used. In other words, the Experiment knows what you would like to run.

class esrgan.core.experiment.SRExperiment(config: Dict)[source]

Experiment for ESRGAN, please check catalyst docs for more info.

get_datasets(stage: str, train_dataset_params: Optional[Dict] = None, valid_dataset_params: Optional[Dict] = None, infer_dataset_params: Optional[Dict] = None) → dict[source]

Returns the datasets for a given stage and epoch.

Parameters

  • stage – stage name of interest, e.g. “train”, “finetune”, “gan” …

  • train_dataset_params – Parameters of train dataset, must contain 'dataset' key with the name of dataset to use e.g. esrgan.dataset.DIV2KDataset.

  • valid_dataset_params – Parameters of validation dataset.

  • infer_dataset_params – Parameters of inference dataset.

Returns

Dictionary with datasets for current stage.

Example for Config API:

train_dataset_params:
  dataset: DIV2KDataset
  root: data
  train: true
  target_type: bicubic_X4
  download: true

GAN Runners

Runner is an abstraction that knows how to run an experiment. It contains all the logic of how to run the experiment, stages, epoch and batches.

class esrgan.core.runner.GANRunner(model: Union[torch.nn.modules.module.Module, Dict[str, torch.nn.modules.module.Module]] = None, device: Union[str, torch.device] = None, **kwargs)[source]

Runner for experiments with supervised / GAN model.

predict_batch(batch: Dict[str, torch.Tensor]) → torch.Tensor[source]

Generate predictions based on input batch (generator inference).

Parameters

batch – Input batch (batch of samples to adjust e.g. zoom).

Returns

Batch of predictions of the generator.

Callbacks

Callback is an abstraction that lets you customize your experiment run logic. To give users maximum flexibility and extensibility Catalyst supports callback execution anywhere in the training loop:

-- stage start
---- epoch start
------ loader start
-------- batch start
---------- batch handler (Runner logic)
-------- batch end
------ loader end
---- epoch end
-- stage end

exception – if an Exception was raised

For example, to calculate ROC-AUC of the model you may use on_batch_end() method to gather per-batch predictions and on_loader_end() method to average those statistics.

Metrics

class esrgan.callbacks.metrics.PSNRCallback(input_key: str = 'targets', output_key: str = 'outputs', prefix: str = 'psnr', multiplier: float = 1.0, data_range: Union[int, float] = 1.0, reduction: str = 'mean', convert_to_greyscale: bool = False)[source]

Peak signal-to-noise ratio (PSNR) metric callback.

Compute Peak Signal-to-Noise Ratio for a batch of images.

Parameters

  • input_key – Input key to use for PSNR calculation; specifies our y_true.

  • output_key – Output key to use for PSNR calculation; specifies our y_pred.

  • prefix – Name of the metric / key to store in logs.

  • multiplier – Scale factor for the metric.

  • data_range – Value range of input images (usually 1.0 or 255).

  • reduction – Reduction over samples in batch, should be one of: 'mean', 'sum', or 'none'.

  • convert_to_greyscale – If True, convert RGB image to YCbCr format and computes PSNR only on luminance channel, compute on all 3 channels otherwise.

class esrgan.callbacks.metrics.SSIMCallback(input_key: str = 'targets', output_key: str = 'outputs', prefix: str = 'ssim', multiplier: float = 1.0, kernel_size: int = 11, kernel_sigma: float = 1.5, data_range: Union[int, float] = 1.0, reduction: str = 'mean', k1: float = 0.01, k2: float = 0.03)[source]

Structural similarity (SSIM) metric callback.

Computes Structural Similarity (SSIM) index between two images. It has been proposed in Image Quality Assessment: From Error Visibility to Structural Similarity.

Parameters

  • input_key – Input key to use for SSIM calculation; specifies our y_true.

  • output_key – Output key to use for SSIM calculation; specifies our y_pred.

  • prefix – Name of the metric / key to store in logs.

  • multiplier – Scale factor for the metric.

  • kernel_size – The side-length of the Gaussian sliding window used in comparison. Must be an odd value.

  • kernel_sigma – Standard deviation of normal distribution.

  • data_range – Value range of input images (usually 1.0 or 255).

  • reduction – Specifies the reduction to apply to the output, should be one of: 'mean', 'sum', or 'none'.

  • k1 – Algorithm parameter, small constant used to stabilize the division with small denominator (see original paper for more info).

  • k2 – Algorithm parameter, small constant used to stabilize the division with small denominator (see original paper for more info).