You are reading the documentation for MMEditing 0.x, which will soon be deprecated by the end of 2022. We recommend you upgrade to MMEditing 1.0 to enjoy fruitful new features and better performance brought by OpenMMLab 2.0. Check out the changelog, code and documentation of MMEditing 1.0 for more details.

Source code for mmedit.models.components.discriminators.multi_layer_disc

# Copyright (c) OpenMMLab. All rights reserved.
import torch.nn as nn
from mmcv.cnn import ConvModule
from mmcv.runner import load_checkpoint

from mmedit.models.common import LinearModule
from mmedit.models.registry import COMPONENTS
from mmedit.utils import get_root_logger

[docs]@COMPONENTS.register_module() class MultiLayerDiscriminator(nn.Module): """Multilayer Discriminator. This is a commonly used structure with stacked multiply convolution layers. Args: in_channels (int): Input channel of the first input convolution. max_channels (int): The maximum channel number in this structure. num_conv (int): Number of stacked intermediate convs (including input conv but excluding output conv). fc_in_channels (int | None): Input dimension of the fully connected layer. If `fc_in_channels` is None, the fully connected layer will be removed. fc_out_channels (int): Output dimension of the fully connected layer. kernel_size (int): Kernel size of the conv modules. Default to 5. conv_cfg (dict): Config dict to build conv layer. norm_cfg (dict): Config dict to build norm layer. act_cfg (dict): Config dict for activation layer, "relu" by default. out_act_cfg (dict): Config dict for output activation, "relu" by default. with_input_norm (bool): Whether add normalization after the input conv. Default to True. with_out_convs (bool): Whether add output convs to the discriminator. The output convs contain two convs. The first out conv has the same setting as the intermediate convs but a stride of 1 instead of 2. The second out conv is a conv similar to the first out conv but reduces the number of channels to 1 and has no activation layer. Default to False. with_spectral_norm (bool): Whether use spectral norm after the conv layers. Default to False. kwargs (keyword arguments). """ def __init__(self, in_channels, max_channels, num_convs=5, fc_in_channels=None, fc_out_channels=1024, kernel_size=5, conv_cfg=None, norm_cfg=None, act_cfg=dict(type='ReLU'), out_act_cfg=dict(type='ReLU'), with_input_norm=True, with_out_convs=False, with_spectral_norm=False, **kwargs): super().__init__() if fc_in_channels is not None: assert fc_in_channels > 0 self.max_channels = max_channels self.with_fc = fc_in_channels is not None self.num_convs = num_convs self.with_out_act = out_act_cfg is not None self.with_out_convs = with_out_convs cur_channels = in_channels for i in range(num_convs): out_ch = min(64 * 2**i, max_channels) norm_cfg_ = norm_cfg act_cfg_ = act_cfg if i == 0 and not with_input_norm: norm_cfg_ = None elif (i == num_convs - 1 and not self.with_fc and not self.with_out_convs): norm_cfg_ = None act_cfg_ = out_act_cfg self.add_module( f'conv{i + 1}', ConvModule( cur_channels, out_ch, kernel_size=kernel_size, stride=2, padding=kernel_size // 2, norm_cfg=norm_cfg_, act_cfg=act_cfg_, with_spectral_norm=with_spectral_norm, **kwargs)) cur_channels = out_ch if self.with_out_convs: cur_channels = min(64 * 2**(num_convs - 1), max_channels) out_ch = min(64 * 2**num_convs, max_channels) self.add_module( f'conv{num_convs + 1}', ConvModule( cur_channels, out_ch, kernel_size, stride=1, padding=kernel_size // 2, norm_cfg=norm_cfg, act_cfg=act_cfg, with_spectral_norm=with_spectral_norm, **kwargs)) self.add_module( f'conv{num_convs + 2}', ConvModule( out_ch, 1, kernel_size, stride=1, padding=kernel_size // 2, act_cfg=None, with_spectral_norm=with_spectral_norm, **kwargs)) if self.with_fc: self.fc = LinearModule( fc_in_channels, fc_out_channels, bias=True, act_cfg=out_act_cfg, with_spectral_norm=with_spectral_norm)
[docs] def forward(self, x): """Forward Function. Args: x (torch.Tensor): Input tensor with shape of (n, c, h, w). Returns: torch.Tensor: Output tensor with shape of (n, c, h', w') or (n, c). """ input_size = x.size() # out_convs has two additional ConvModules num_convs = self.num_convs + 2 * self.with_out_convs for i in range(num_convs): x = getattr(self, f'conv{i + 1}')(x) if self.with_fc: x = x.view(input_size[0], -1) x = self.fc(x) return x
[docs] def init_weights(self, pretrained=None): """Init weights for models. Args: pretrained (str, optional): Path for pretrained weights. If given None, pretrained weights will not be loaded. Defaults to None. """ if isinstance(pretrained, str): logger = get_root_logger() load_checkpoint(self, pretrained, strict=False, logger=logger) elif pretrained is None: for m in self.modules(): # Here, we only initialize the module with fc layer since the # conv and norm layers has been initialized in `ConvModule`. if isinstance(m, nn.Linear): nn.init.normal_(, 0.0, 0.02) nn.init.constant_(, 0.0) else: raise TypeError('pretrained must be a str or None')
Read the Docs v: latest
On Read the Docs
Project Home

Free document hosting provided by Read the Docs.