Catalyst

Various features for customization/modification of Catalyst pipelines e.g., runners, metrics, callbacks:

Runners

GANRunner

class esrgan.runner.GANRunner(input_key: str = 'image', target_key: str = 'real_image', generator_output_key: str = 'fake_image', discriminator_real_output_gkey: str = 'g_real_logits', discriminator_fake_output_gkey: str = 'g_fake_logits', discriminator_real_output_dkey: str = 'd_real_logits', discriminator_fake_output_dkey: str = 'd_fake_logits', generator_key: str = 'generator', discriminator_key: str = 'discriminator')

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

Parameters

• input_key – Key in batch dict mapping for model input.

• target_key – Key in batch dict mapping for target.

• generator_output_key – Key in output dict model output of the generator will be stored under.

• discriminator_real_output_gkey – Key to store predictions of discriminator for real inputs, contain gradients for generator.

• discriminator_fake_output_gkey – Key to store predictions of discriminator for predictions of generator, contain gradients for generator.

• discriminator_real_output_dkey – Key to store predictions of discriminator for real inputs, contain gradients for discriminator only.

• discriminator_fake_output_dkey – Key to store predictions of discriminator for items produced by generator, contain gradients for discriminator only.

• generator_key – Key in model dict mapping for generator model.

• discriminator_key – Key in model dict mapping for discriminator model (will be used in gan stages only).

handle_batch(batch: Dict[str, torch.Tensor]) → None

Inner method to handle specified data batch.

Parameters

batch – Dictionary with data batches from DataLoader.

on_stage_start(runner: catalyst.core.runner.IRunner) → None

Prepare _handle_batch method for current stage.

Parameters

runner – Current runner.

Raises

NotImplementedError – Name of the stage should ends with '_supervised', '_gan' or should be 'infer', raise error otherwise.

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

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.

GANConfigRunner

class esrgan.runner.GANConfigRunner(config: Dict, input_key: str = 'image', target_key: str = 'real_image', generator_output_key: str = 'fake_image', discriminator_real_output_gkey: str = 'g_real_logits', discriminator_fake_output_gkey: str = 'g_fake_logits', discriminator_real_output_dkey: str = 'd_real_logits', discriminator_fake_output_dkey: str = 'd_fake_logits', generator_key: str = 'generator', discriminator_key: str = 'discriminator')

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

Parameters

• config – Dictionary with parameters e.g., model or engine to use.

• input_key – Key in batch dict mapping for model input.

• target_key – Key in batch dict mapping for target.

• generator_output_key – Key in output dict model output of the generator will be stored under.

• discriminator_real_output_gkey – Key to store predictions of discriminator for real inputs, contain gradients for generator.

• discriminator_fake_output_gkey – Key to store predictions of discriminator for predictions of generator, contain gradients for generator.

• discriminator_real_output_dkey – Key to store predictions of discriminator for real inputs, contain gradients for discriminator only.

• discriminator_fake_output_dkey – Key to store predictions of discriminator for items produced by generator, contain gradients for discriminator only.

• generator_key – Key in model dict mapping for generator model.

• discriminator_key – Key in model dict mapping for discriminator model (will be used in gan stages only).