mmagic.models.editors.pix2pix

Package Contents

Classes

Pix2Pix	Pix2Pix model for paired image-to-image translation.
UnetGenerator	Construct the Unet-based generator from the innermost layer to the

class mmagic.models.editors.pix2pix.Pix2Pix(*args, **kwargs)

Bases: mmagic.models.base_models.BaseTranslationModel

Pix2Pix model for paired image-to-image translation.

Ref:

Image-to-Image Translation with Conditional Adversarial Networks

forward_test(img, target_domain, **kwargs)

Forward function for testing.

Parameters

• img (tensor) – Input image tensor.

• target_domain (str) – Target domain of output image.

• kwargs (dict) – Other arguments.

Returns

Forward results.

Return type

dict

_get_disc_loss(outputs)

Get the loss of discriminator.

Parameters

outputs (dict) – A dict of output.

Returns

Loss and a dict of log of loss terms.

Return type

Tuple

_get_gen_loss(outputs)

Get the loss of generator.

Parameters

outputs (dict) – A dict of output.

Returns

Loss and a dict of log of loss terms.

Return type

Tuple

train_step(data, optim_wrapper=None)

Training step function.

Parameters

• data_batch (dict) – Dict of the input data batch.

• optimizer (dict[torch.optim.Optimizer]) – Dict of optimizers for the generator and discriminator.

• ddp_reducer (Reducer | None, optional) – Reducer from ddp. It is used to prepare for backward() in ddp. Defaults to None.

• running_status (dict | None, optional) – Contains necessary basic information for training, e.g., iteration number. Defaults to None.

Returns

Dict of loss, information for logger, the number of samples and results for visualization.

Return type

dict

test_step(data: dict) → mmagic.utils.typing.SampleList

Gets the generated image of given data. Same as val_step().

Parameters

data (dict) – Data sampled from metric specific sampler. More details in Metrics and Evaluator.

Returns

Generated image or image dict.

Return type

List[DataSample]

val_step(data: dict) → mmagic.utils.typing.SampleList

Gets the generated image of given data. Same as val_step().

Parameters

data (dict) – Data sampled from metric specific sampler. More details in Metrics and Evaluator.

Returns

Generated image or image dict.

Return type

List[DataSample]

class mmagic.models.editors.pix2pix.UnetGenerator(in_channels, out_channels, num_down=8, base_channels=64, norm_cfg=dict(type='BN'), use_dropout=False, init_cfg=dict(type='normal', gain=0.02))

Bases: mmengine.model.BaseModule

Construct the Unet-based generator from the innermost layer to the outermost layer, which is a recursive process.

Parameters

• in_channels (int) – Number of channels in input images.

• out_channels (int) – Number of channels in output images.

• num_down (int) – Number of downsamplings in Unet. If num_down is 8, the image with size 256x256 will become 1x1 at the bottleneck. Default: 8.

• base_channels (int) – Number of channels at the last conv layer. Default: 64.

• norm_cfg (dict) – Config dict to build norm layer. Default: dict(type=’BN’).

• use_dropout (bool) – Whether to use dropout layers. Default: False.

• init_cfg (dict) – Config dict for initialization. type: The name of our initialization method. Default: ‘normal’. gain: Scaling factor for normal, xavier and orthogonal. Default: 0.02.

forward(x)

Forward function.

Parameters

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

Returns

Forward results.

Return type

Tensor

init_weights()

Initialize weights for the model.

Parameters

• pretrained (str, optional) – Path for pretrained weights. If given None, pretrained weights will not be loaded. Default: None.

• strict (bool, optional) – Whether to allow different params for the model and checkpoint. Default: True.