mmagic.models.editors.esrgan

Package Contents

Classes

ESRGAN	Enhanced SRGAN model for single image super-resolution.
RRDBNet	Networks consisting of Residual in Residual Dense Block, which is used

class mmagic.models.editors.esrgan.ESRGAN(generator, discriminator=None, gan_loss=None, pixel_loss=None, perceptual_loss=None, train_cfg=None, test_cfg=None, init_cfg=None, data_preprocessor=None)

Bases: mmagic.models.editors.srgan.SRGAN

Enhanced SRGAN model for single image super-resolution.

Ref: ESRGAN: Enhanced Super-Resolution Generative Adversarial Networks. It uses RaGAN for GAN updates: The relativistic discriminator: a key element missing from standard GAN.

Parameters

• generator (dict) – Config for the generator.

• discriminator (dict) – Config for the discriminator. Default: None.

• gan_loss (dict) – Config for the gan loss. Note that the loss weight in gan loss is only for the generator.

• pixel_loss (dict) – Config for the pixel loss. Default: None.

• perceptual_loss (dict) – Config for the perceptual loss. Default: None.

• train_cfg (dict) – Config for training. Default: None. You may change the training of gan by setting: disc_steps: how many discriminator updates after one generate update; disc_init_steps: how many discriminator updates at the start of the training. These two keys are useful when training with WGAN.

• test_cfg (dict) – Config for testing. Default: None.

• init_cfg (dict, optional) – The weight initialized config for BaseModule. Default: None.

g_step(batch_outputs: torch.Tensor, batch_gt_data: torch.Tensor)

G step of GAN: Calculate losses of generator.

Parameters

• batch_outputs (Tensor) – Batch output of generator.

• batch_gt_data (Tensor) – Batch GT data.

Returns

Dict of losses.

Return type

dict

d_step_real(batch_outputs: torch.Tensor, batch_gt_data: torch.Tensor)

D step of real data.

Parameters

• batch_outputs (Tensor) – Batch output of generator.

• batch_gt_data (Tensor) – Batch GT data.

Returns

Dict of losses.

Return type

dict

d_step_fake(batch_outputs: torch.Tensor, batch_gt_data)

D step of fake data.

Parameters

• batch_outputs (Tensor) – Batch output of generator.

• batch_gt_data (Tensor) – Batch GT data.

Returns

Dict of losses.

Return type

dict

class mmagic.models.editors.esrgan.RRDBNet(in_channels, out_channels, mid_channels=64, num_blocks=23, growth_channels=32, upscale_factor=4, init_cfg=None)

Bases: mmengine.model.BaseModule

Networks consisting of Residual in Residual Dense Block, which is used in ESRGAN and Real-ESRGAN.

ESRGAN: Enhanced Super-Resolution Generative Adversarial Networks. Real-ESRGAN: Training Real-World Blind Super-Resolution with Pure Synthetic Data. # noqa: E501 Currently, it supports [x1/x2/x4] upsampling scale factor.

Parameters

• in_channels (int) – Channel number of inputs.

• out_channels (int) – Channel number of outputs.

• mid_channels (int) – Channel number of intermediate features. Default: 64

• num_blocks (int) – Block number in the trunk network. Defaults: 23

• growth_channels (int) – Channels for each growth. Default: 32.

• upscale_factor (int) – Upsampling factor. Support x1, x2 and x4. Default: 4.

• init_cfg (dict, optional) – Initialization config dict. Default: None.

_supported_upscale_factors = [1, 2, 4]

forward(x)

Forward function.

Parameters

x (Tensor) – Input tensor with shape (n, c, h, w).

Returns

Forward results.

Return type

Tensor

init_weights()

Init weights for models.