mmedit.models.editors.biggan.biggan_discriminator

Module Contents

Classes

BigGANDiscriminator

BigGAN Discriminator. The implementation refers to

class mmedit.models.editors.biggan.biggan_discriminator.BigGANDiscriminator(input_scale, num_classes=0, in_channels=3, out_channels=1, base_channels=96, sn_eps=1e-06, sn_style='ajbrock', init_type='ortho', act_cfg=dict(type='ReLU'), with_spectral_norm=True, blocks_cfg=dict(type='BigGANDiscResBlock'), arch_cfg=None, pretrained=None)

Bases: torch.nn.Module

BigGAN Discriminator. The implementation refers to https://github.com/ajbrock/BigGAN-PyTorch/blob/master/BigGAN.py # noqa.

In BigGAN, we use a SAGAN-based architecture composing of an self-attention block and number of convolutional residual blocks with spectral normalization.

More details can be found in: Large Scale GAN Training for High Fidelity Natural Image Synthesis (ICLR2019).

The design of the model structure is highly corresponding to the output resolution. For the original BigGAN’s generator, you can set output_scale as you need and use the default value of arch_cfg and blocks_cfg. If you want to customize the model, you can set the arguments in this way:

arch_cfg: Config for the architecture of this generator. You can refer the _default_arch_cfgs in the _get_default_arch_cfg function to see the format of the arch_cfg. Basically, you need to provide information of each block such as the numbers of input and output channels, whether to perform upsampling, etc.

blocks_cfg: Config for the convolution block. You can replace the block type to your registered customized block and adjust block params here. However, you should notice that some params are shared among these blocks like act_cfg, with_spectral_norm, sn_eps, etc.

参数

• input_scale (int) – The scale of the input image.

• num_classes (int, optional) – The number of conditional classes. Defaults to 0.

• in_channels (int, optional) – The channel number of the input image. Defaults to 3.

• out_channels (int, optional) – The channel number of the final output. Defaults to 1.

• base_channels (int, optional) – The basic channel number of the discriminator. The other layers contains channels based on this number. Defaults to 96.

• sn_eps (float, optional) – Epsilon value for spectral normalization. Defaults to 1e-6.

• sn_style (str, optional) – The style of spectral normalization. If set to ajbrock, implementation by ajbrock(https://github.com/ajbrock/BigGAN-PyTorch/blob/master/layers.py) will be adopted. If set to torch, implementation by PyTorch will be adopted. Defaults to ajbrock.

• init_type (str, optional) – The name of an initialization method: ortho | N02 | xavier. Defaults to ‘ortho’.

• act_cfg (dict, optional) – Config for the activation layer. Defaults to dict(type=’ReLU’).

• with_spectral_norm (bool, optional) – Whether to use spectral normalization. Defaults to True.

• blocks_cfg (dict, optional) – Config for the convolution block. Defaults to dict(type=’BigGANDiscResBlock’).

• arch_cfg (dict, optional) – Config for the architecture of this discriminator. Defaults to None.

• pretrained (str | dict, optional) – Path for the pretrained model or dict containing information for pretained models whose necessary key is ‘ckpt_path’. Besides, you can also provide ‘prefix’ to load the generator part from the whole state dict. Defaults to None.

_get_default_arch_cfg(input_scale, in_channels, base_channels)

forward(x, label=None)

Forward function.

参数

• x (torch.Tensor) – Fake or real image tensor.

• label (torch.Tensor | None) – Label Tensor. Defaults to None.

返回

Prediction for the reality of the input image with

given label.

返回类型

torch.Tensor

init_weights(pretrained=None, init_type='ortho')

Init weights for models.

参数

• pretrained (str | dict, optional) – Path for the pretrained model or dict containing information for pretained models whose necessary key is ‘ckpt_path’. Besides, you can also provide ‘prefix’ to load the generator part from the whole state dict. Defaults to None.

• init_type (str, optional) – The name of an initialization method: ortho | N02 | xavier. Defaults to ‘ortho’.