mmagic.models.editors.pconv.pconv_encoder

Module Contents

Classes

PConvEncoder

Encoder with partial conv.

class mmagic.models.editors.pconv.pconv_encoder.PConvEncoder(in_channels=3, num_layers=7, conv_cfg=dict(type='PConv', multi_channel=True), norm_cfg=dict(type='BN', requires_grad=True), norm_eval=False)

Bases: mmengine.model.BaseModule

Encoder with partial conv.

About the details for this architecture, pls see: Image Inpainting for Irregular Holes Using Partial Convolutions

Parameters

• in_channels (int) – The number of input channels. Default: 3.

• num_layers (int) – The number of convolutional layers. Default: 7.

• conv_cfg (dict) – Config for convolution module. Default: {‘type’: ‘PConv’, ‘multi_channel’: True}.

• norm_cfg (dict) – Config for norm layer. Default: {‘type’: ‘BN’}.

• norm_eval (bool) – Whether to set norm layers to eval mode, namely, freeze running stats (mean and var). Note: Effective on Batch Norm and its variants only. Default: False.

train(mode=True)

Set BatchNorm modules in the model to evaluation mode.

forward(x, mask)

Forward function for partial conv encoder.

Parameters

• x (torch.Tensor) – Masked image with shape (n, c, h, w).

• mask (torch.Tensor) – Mask tensor with shape (n, c, h, w).

Returns

Contains the results and middle level features in this module. hidden_feats contain the middle feature maps and hidden_masks store updated masks.

Return type

dict