def test_build_activation_layer(): with pytest.raises(TypeError): # cfg must be a dict cfg = 'ReLU' build_activation_layer(cfg) with pytest.raises(KeyError): # `type` must be in cfg cfg = dict() build_activation_layer(cfg) with pytest.raises(KeyError): # unsupported activation type cfg = dict(type='FancyReLU') build_activation_layer(cfg) # test each type of activation layer in activation_cfg for type_name, module in ACTIVATION_LAYERS.module_dict.items(): cfg['type'] = type_name layer = build_activation_layer(cfg) assert isinstance(layer, module) # sanity check for Clamp act = build_activation_layer(dict(type='Clamp')) x = torch.randn(10) * 1000 y = act(x) assert np.logical_and((y >= -1).numpy(), (y <= 1).numpy()).all() act = build_activation_layer(dict(type='Clip', min=0)) y = act(x) assert np.logical_and((y >= 0).numpy(), (y <= 1).numpy()).all() act = build_activation_layer(dict(type='Clamp', max=0)) y = act(x) assert np.logical_and((y >= -1).numpy(), (y <= 0).numpy()).all()
def __init__(self, in_channels, embedding_channels, use_scale_shift_norm, dropout, out_channels=None, norm_cfg=dict(type='GN', num_groups=32), act_cfg=dict(type='SiLU', inplace=False), shortcut_kernel_size=1): super().__init__() out_channels = in_channels if out_channels is None else out_channels _norm_cfg = deepcopy(norm_cfg) _, norm_1 = build_norm_layer(_norm_cfg, in_channels) conv_1 = [ norm_1, build_activation_layer(act_cfg), nn.Conv2d(in_channels, out_channels, 3, padding=1) ] self.conv_1 = nn.Sequential(*conv_1) norm_with_embedding_cfg = dict( in_channels=out_channels, embedding_channels=embedding_channels, use_scale_shift=use_scale_shift_norm, norm_cfg=_norm_cfg) self.norm_with_embedding = build_module( dict(type='NormWithEmbedding'), default_args=norm_with_embedding_cfg) conv_2 = [ build_activation_layer(act_cfg), nn.Dropout(dropout), nn.Conv2d(out_channels, out_channels, 3, padding=1) ] self.conv_2 = nn.Sequential(*conv_2) assert shortcut_kernel_size in [ 1, 3 ], ('Only support `1` and `3` for `shortcut_kernel_size`, but ' f'receive {shortcut_kernel_size}.') self.learnable_shortcut = out_channels != in_channels if self.learnable_shortcut: shortcut_padding = 1 if shortcut_kernel_size == 3 else 0 self.shortcut = nn.Conv2d( in_channels, out_channels, shortcut_kernel_size, padding=shortcut_padding) self.init_weights()
def __init__(self, in_channels, norm_cfg=dict(type='LN2d', eps=1e-6), act_cfg=dict(type='GELU'), mlp_ratio=4., linear_pw_conv=True, drop_path_rate=0., layer_scale_init_value=1e-6): super().__init__() self.depthwise_conv = nn.Conv2d(in_channels, in_channels, kernel_size=7, padding=3, groups=in_channels) self.linear_pw_conv = linear_pw_conv self.norm = build_norm_layer(norm_cfg, in_channels)[1] mid_channels = int(mlp_ratio * in_channels) if self.linear_pw_conv: # Use linear layer to do pointwise conv. pw_conv = nn.Linear else: pw_conv = partial(nn.Conv2d, kernel_size=1) self.pointwise_conv1 = pw_conv(in_channels, mid_channels) self.act = build_activation_layer(act_cfg) self.pointwise_conv2 = pw_conv(mid_channels, in_channels) self.gamma = nn.Parameter( layer_scale_init_value * torch.ones((in_channels)), requires_grad=True) if layer_scale_init_value > 0 else None self.drop_path = DropPath( drop_path_rate) if drop_path_rate > 0. else nn.Identity()
def __init__(self, noise_size, out_channels, act_cfg=dict(type='LeakyReLU', negative_slope=0.2), norm_cfg=dict(type='PixelNorm'), normalize_latent=True, order=('linear', 'act', 'norm')): super().__init__() self.noise_size = noise_size self.out_channels = out_channels self.normalize_latent = normalize_latent self.with_activation = act_cfg is not None self.with_norm = norm_cfg is not None self.order = order assert len(order) == 3 and set(order) == set(['linear', 'act', 'norm']) # w/o bias, because the bias is added after reshaping the tensor to # 2D feature self.linear = EqualizedLRLinearModule( noise_size, out_channels * 16, equalized_lr_cfg=dict(gain=np.sqrt(2) / 4), bias=False) if self.with_activation: self.activation = build_activation_layer(act_cfg) # add bias for reshaped 2D feature. self.register_parameter( 'bias', nn.Parameter(torch.zeros(1, out_channels, 1, 1))) if self.with_norm: _, self.norm = build_norm_layer(norm_cfg, out_channels)
def __init__(self, in_channels, growth_rate, bn_size, norm_cfg=dict(type='BN'), act_cfg=dict(type='ReLU'), drop_rate=0., memory_efficient=False): super(DenseLayer, self).__init__() self.norm1 = build_norm_layer(norm_cfg, in_channels)[1] self.conv1 = nn.Conv2d(in_channels, bn_size * growth_rate, kernel_size=1, stride=1, bias=False) self.act = build_activation_layer(act_cfg) self.norm2 = build_norm_layer(norm_cfg, bn_size * growth_rate)[1] self.conv2 = nn.Conv2d(bn_size * growth_rate, growth_rate, kernel_size=3, stride=1, padding=1, bias=False) self.drop_rate = float(drop_rate) self.memory_efficient = memory_efficient
def __init__(self, input_scale=128, output_scale=8, out_channels=1, in_channels=3, base_channels=64, conv_cfg=dict(type='Conv2d'), default_norm_cfg=dict(type='BN'), default_act_cfg=dict(type='LeakyReLU', negative_slope=0.2), out_act_cfg=None): super().__init__() assert input_scale % output_scale == 0 assert input_scale // output_scale >= 2 self.input_scale = input_scale self.output_scale = output_scale self.out_channels = out_channels self.base_channels = base_channels self.with_out_activation = out_act_cfg is not None self.conv_blocks = nn.ModuleList() self.conv_blocks.append( ConvModule( in_channels, base_channels, kernel_size=5, stride=2, padding=2, conv_cfg=conv_cfg, norm_cfg=None, act_cfg=default_act_cfg)) # the number of times for downsampling self.num_downsamples = int(np.log2(input_scale // output_scale)) - 1 # build up downsampling backbone (excluding the output layer) curr_channels = base_channels for _ in range(self.num_downsamples): self.conv_blocks.append( ConvModule( curr_channels, curr_channels * 2, kernel_size=5, stride=2, padding=2, conv_cfg=conv_cfg, norm_cfg=default_norm_cfg, act_cfg=default_act_cfg)) curr_channels = curr_channels * 2 # output layer self.decision = nn.Sequential( nn.Linear(output_scale * output_scale * curr_channels, out_channels)) if self.with_out_activation: self.out_activation = build_activation_layer(out_act_cfg)
def __init__(self, input_scale, num_classes=0, in_channels=3, out_channels=1, base_channels=96, sn_eps=1e-6, init_type='ortho', act_cfg=dict(type='ReLU'), with_spectral_norm=True, blocks_cfg=dict(type='BigGANDiscResBlock'), arch_cfg=None, pretrained=None): super().__init__() self.num_classes = num_classes self.out_channels = out_channels self.input_scale = input_scale self.in_channels = in_channels self.base_channels = base_channels self.arch = arch_cfg if arch_cfg else self._get_default_arch_cfg( self.input_scale, self.in_channels, self.base_channels) self.blocks_cfg = deepcopy(blocks_cfg) self.blocks_cfg.update( dict(act_cfg=act_cfg, sn_eps=sn_eps, with_spectral_norm=with_spectral_norm)) self.conv_blocks = nn.ModuleList() for index, out_ch in enumerate(self.arch['out_channels']): # change args to adapt to current block self.blocks_cfg.update( dict(in_channels=self.arch['in_channels'][index], out_channels=out_ch, with_downsample=self.arch['downsample'][index], is_head_block=(index == 0))) self.conv_blocks.append(build_module(self.blocks_cfg)) if self.arch['attention'][index]: self.conv_blocks.append( SelfAttentionBlock(out_ch, with_spectral_norm=with_spectral_norm, sn_eps=sn_eps)) self.activate = build_activation_layer(act_cfg) self.decision = nn.Linear(self.arch['out_channels'][-1], out_channels) if with_spectral_norm: self.decision = spectral_norm(self.decision, eps=sn_eps) if self.num_classes > 0: self.proj_y = nn.Embedding(self.num_classes, self.arch['out_channels'][-1]) if with_spectral_norm: self.proj_y = spectral_norm(self.proj_y, eps=sn_eps) self.init_weights(pretrained=pretrained, init_type=init_type)
def __init__(self, in_features, hidden_features=None, out_features=None, act_cfg=dict(type='GELU'), drop=0.): super().__init__() out_features = out_features or in_features hidden_features = hidden_features or in_features self.fc1 = nn.Conv2d(in_features, hidden_features, 1) self.act = build_activation_layer(act_cfg) self.fc2 = nn.Conv2d(hidden_features, out_features, 1) self.drop = nn.Dropout(drop)
def __init__(self, in_channels, mid_channels, out_channels, bias=True, equalized_lr_cfg=dict(gain=1), act_cfg=dict(type='LeakyReLU', negative_slope=0.2), out_act=None): super().__init__() self.in_channels = in_channels self.mid_channels = mid_channels self.out_channels = out_channels self.with_activation = act_cfg is not None self.with_out_activation = out_act is not None # setup linear layers # dirty code for supporting default mode in PGGAN if equalized_lr_cfg: equalized_lr_cfg_ = dict(gain=2**0.5) else: equalized_lr_cfg_ = None self.linear0 = EqualizedLRLinearModule( self.in_channels, self.mid_channels, bias=bias, equalized_lr_cfg=equalized_lr_cfg_) self.linear1 = EqualizedLRLinearModule( self.mid_channels, self.out_channels, bias=bias, equalized_lr_cfg=equalized_lr_cfg) # setup activation layers if self.with_activation: self.activation = build_activation_layer(act_cfg) if self.with_out_activation: self.out_activation = build_activation_layer(out_act)
def __init__(self, in_channels, embedding_channels, norm_cfg=dict(type='GN', num_groups=32), act_cfg=dict(type='SiLU', inplace=False), use_scale_shift=True): super().__init__() self.use_scale_shift = use_scale_shift _, self.norm = build_norm_layer(norm_cfg, in_channels) embedding_output = in_channels * 2 if use_scale_shift else in_channels self.embedding_layer = nn.Sequential( build_activation_layer(act_cfg), nn.Linear(embedding_channels, embedding_output))
def __init__(self, in_channels, out_channels, norm_cfg=dict(type='BN'), act_cfg=dict(type='ReLU')): super(DenseTransition, self).__init__() self.add_module('norm', build_norm_layer(norm_cfg, in_channels)[1]) self.add_module('act', build_activation_layer(act_cfg)) self.add_module( 'conv', nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, bias=False)) self.add_module('pool', nn.AvgPool2d(kernel_size=2, stride=2))
def __init__(self, in_channels, out_channels, act_cfg=dict(type='ReLU', inplace=False), sn_eps=1e-6, sn_style='ajbrock', with_downsample=True, with_spectral_norm=True, is_head_block=False): super().__init__() self.activation = build_activation_layer(act_cfg) self.with_downsample = with_downsample self.is_head_block = is_head_block if self.with_downsample: self.downsample = nn.AvgPool2d(kernel_size=2, stride=2) self.learnable_sc = in_channels != out_channels or self.with_downsample if self.learnable_sc: self.shortcut = SNConvModule(in_channels=in_channels, out_channels=out_channels, kernel_size=1, stride=1, padding=0, act_cfg=None, with_spectral_norm=with_spectral_norm, spectral_norm_cfg=dict( eps=sn_eps, sn_style=sn_style)) self.conv1 = SNConvModule(in_channels=in_channels, out_channels=out_channels, kernel_size=3, stride=1, padding=1, act_cfg=None, with_spectral_norm=with_spectral_norm, spectral_norm_cfg=dict(eps=sn_eps, sn_style=sn_style)) self.conv2 = SNConvModule(in_channels=out_channels, out_channels=out_channels, kernel_size=3, stride=1, padding=1, act_cfg=None, with_spectral_norm=with_spectral_norm, spectral_norm_cfg=dict(eps=sn_eps, sn_style=sn_style))
def __init__(self, in_channels, embedding_channels, embedding_mode='sin', embedding_cfg=None, act_cfg=dict(type='SiLU', inplace=False)): super().__init__() self.blocks = nn.Sequential( nn.Linear(in_channels, embedding_channels), build_activation_layer(act_cfg), nn.Linear(embedding_channels, embedding_channels)) # add `dim` to embedding config embedding_cfg_ = dict(dim=in_channels) if embedding_cfg is not None: embedding_cfg_.update(embedding_cfg) if embedding_mode.upper() == 'SIN': self.embedding_fn = partial(self.sinusodial_embedding, **embedding_cfg_) else: raise ValueError('Only support `SIN` for time embedding, ' f'but receive {embedding_mode}.')
def test_build_activation_layer(): with pytest.raises(TypeError): # cfg must be a dict cfg = 'ReLU' build_activation_layer(cfg) with pytest.raises(KeyError): # `type` must be in cfg cfg = dict() build_activation_layer(cfg) with pytest.raises(KeyError): # unsupported activation type cfg = dict(type='FancyReLU') build_activation_layer(cfg) # test each type of activation layer in activation_cfg for type_name, module in ACTIVATION_LAYERS.module_dict.items(): cfg['type'] = type_name layer = build_activation_layer(cfg) assert isinstance(layer, module)
def __init__(self, output_scale=128, out_channels=3, base_channels=256, input_scale=8, noise_size=1024, conv_cfg=dict(type='ConvTranspose2d'), default_norm_cfg=dict(type='BN'), default_act_cfg=dict(type='ReLU'), out_act_cfg=dict(type='Tanh')): super().__init__() assert output_scale % input_scale == 0 assert output_scale // input_scale >= 4 self.output_scale = output_scale self.base_channels = base_channels self.input_scale = input_scale self.noise_size = noise_size self.noise2feat_head = nn.Sequential( nn.Linear(noise_size, input_scale * input_scale * base_channels)) self.noise2feat_tail = nn.Sequential(nn.BatchNorm2d(base_channels)) if default_act_cfg is not None: self.noise2feat_tail.add_module( 'act', build_activation_layer(default_act_cfg)) # the number of times for upsampling self.num_upsamples = int(np.log2(output_scale // input_scale)) - 2 # build up convolution backbone (excluding the output layer) self.conv_blocks = nn.ModuleList() for _ in range(self.num_upsamples): self.conv_blocks.append( ConvModule(base_channels, base_channels, kernel_size=3, stride=2, padding=1, conv_cfg=dict(conv_cfg, output_padding=1), norm_cfg=default_norm_cfg, act_cfg=default_act_cfg)) self.conv_blocks.append( ConvModule(base_channels, base_channels, kernel_size=3, stride=1, padding=1, conv_cfg=conv_cfg, norm_cfg=default_norm_cfg, act_cfg=default_act_cfg)) # output blocks self.conv_blocks.append( ConvModule(base_channels, int(base_channels // 2), kernel_size=3, stride=2, padding=1, conv_cfg=dict(conv_cfg, output_padding=1), norm_cfg=default_norm_cfg, act_cfg=default_act_cfg)) self.conv_blocks.append( ConvModule(int(base_channels // 2), int(base_channels // 4), kernel_size=3, stride=2, padding=1, conv_cfg=dict(conv_cfg, output_padding=1), norm_cfg=default_norm_cfg, act_cfg=default_act_cfg)) self.conv_blocks.append( ConvModule(int(base_channels // 4), out_channels, kernel_size=3, stride=1, padding=1, conv_cfg=conv_cfg, norm_cfg=None, act_cfg=out_act_cfg))
def __init__(self, in_channels, out_channels, dim_after_concat, act_cfg=dict(type='ReLU'), upsample_cfg=dict(type='nearest', scale_factor=2), sn_eps=1e-6, with_spectral_norm=True, input_is_label=False, auto_sync_bn=True): super().__init__() self.activation = build_activation_layer(act_cfg) self.upsample_cfg = deepcopy(upsample_cfg) self.with_upsample = upsample_cfg is not None if self.with_upsample: self.upsample_layer = build_upsample_layer(self.upsample_cfg) self.learnable_sc = in_channels != out_channels or self.with_upsample if self.learnable_sc: self.shortcut = SNConvModule( in_channels=in_channels, out_channels=out_channels, kernel_size=1, stride=1, padding=0, act_cfg=None, with_spectral_norm=with_spectral_norm, spectral_norm_cfg=dict(eps=sn_eps)) # Here in_channels of BigGANGenResBlock equal to num_features of # BigGANConditionBN self.bn1 = BigGANConditionBN( in_channels, dim_after_concat, sn_eps=sn_eps, input_is_label=input_is_label, with_spectral_norm=with_spectral_norm, auto_sync_bn=auto_sync_bn) # Here out_channels of BigGANGenResBlock equal to num_features of # BigGANConditionBN self.bn2 = BigGANConditionBN( out_channels, dim_after_concat, sn_eps=sn_eps, input_is_label=input_is_label, with_spectral_norm=with_spectral_norm, auto_sync_bn=auto_sync_bn) self.conv1 = SNConvModule( in_channels=in_channels, out_channels=out_channels, kernel_size=3, stride=1, padding=1, act_cfg=None, with_spectral_norm=with_spectral_norm, spectral_norm_cfg=dict(eps=sn_eps)) self.conv2 = SNConvModule( in_channels=out_channels, out_channels=out_channels, kernel_size=3, stride=1, padding=1, act_cfg=None, with_spectral_norm=with_spectral_norm, spectral_norm_cfg=dict(eps=sn_eps))
def __init__(self, in_channels, out_channels, channel_ratio=4, act_cfg=dict(type='ReLU', inplace=False), sn_eps=1e-6, with_downsample=True, with_spectral_norm=True): super().__init__() self.in_channels = in_channels self.out_channels = out_channels self.hidden_channels = self.out_channels // channel_ratio self.activation = build_activation_layer(act_cfg) self.with_downsample = with_downsample if self.with_downsample: self.downsample = nn.AvgPool2d(kernel_size=2, stride=2) self.learnable_sc = (in_channels != out_channels) if self.learnable_sc: self.shortcut = SNConvModule( in_channels=in_channels, out_channels=out_channels - in_channels, kernel_size=1, stride=1, padding=0, act_cfg=None, with_spectral_norm=with_spectral_norm, spectral_norm_cfg=dict(eps=sn_eps)) self.conv1 = SNConvModule( in_channels=in_channels, out_channels=self.hidden_channels, kernel_size=1, stride=1, padding=0, act_cfg=act_cfg, with_spectral_norm=with_spectral_norm, spectral_norm_cfg=dict(eps=sn_eps), order=('act', 'conv', 'norm')) self.conv2 = SNConvModule( in_channels=self.hidden_channels, out_channels=self.hidden_channels, kernel_size=3, stride=1, padding=1, act_cfg=act_cfg, with_spectral_norm=with_spectral_norm, spectral_norm_cfg=dict(eps=sn_eps), order=('act', 'conv', 'norm')) self.conv3 = SNConvModule( in_channels=self.hidden_channels, out_channels=self.hidden_channels, kernel_size=3, stride=1, padding=1, act_cfg=act_cfg, with_spectral_norm=with_spectral_norm, spectral_norm_cfg=dict(eps=sn_eps), order=('act', 'conv', 'norm')) self.conv4 = SNConvModule( in_channels=self.hidden_channels, out_channels=out_channels, kernel_size=1, stride=1, padding=0, act_cfg=None, with_spectral_norm=with_spectral_norm, spectral_norm_cfg=dict(eps=sn_eps))
def __init__(self, arch='121', in_channels=3, bn_size=4, drop_rate=0, compression_factor=0.5, memory_efficient=False, norm_cfg=dict(type='BN'), act_cfg=dict(type='ReLU'), out_indices=-1, frozen_stages=0, init_cfg=None): super().__init__(init_cfg=init_cfg) if isinstance(arch, str): assert arch in self.arch_settings, \ f'Unavailable arch, please choose from ' \ f'({set(self.arch_settings)}) or pass a dict.' arch = self.arch_settings[arch] elif isinstance(arch, dict): essential_keys = {'growth_rate', 'depths', 'init_channels'} assert isinstance(arch, dict) and essential_keys <= set(arch), \ f'Custom arch needs a dict with keys {essential_keys}' self.growth_rate = arch['growth_rate'] self.depths = arch['depths'] self.init_channels = arch['init_channels'] self.act = build_activation_layer(act_cfg) self.num_stages = len(self.depths) # check out indices and frozen stages if isinstance(out_indices, int): out_indices = [out_indices] assert isinstance(out_indices, Sequence), \ f'"out_indices" must by a sequence or int, ' \ f'get {type(out_indices)} instead.' for i, index in enumerate(out_indices): if index < 0: out_indices[i] = self.num_stages + index assert out_indices[i] >= 0, f'Invalid out_indices {index}' self.out_indices = out_indices self.frozen_stages = frozen_stages # Set stem layers self.stem = nn.Sequential( nn.Conv2d(in_channels, self.init_channels, kernel_size=7, stride=2, padding=3, bias=False), build_norm_layer(norm_cfg, self.init_channels)[1], self.act, nn.MaxPool2d(kernel_size=3, stride=2, padding=1)) # Repetitions of DenseNet Blocks self.stages = nn.ModuleList() self.transitions = nn.ModuleList() channels = self.init_channels for i in range(self.num_stages): depth = self.depths[i] stage = DenseBlock(num_layers=depth, in_channels=channels, bn_size=bn_size, growth_rate=self.growth_rate, norm_cfg=norm_cfg, act_cfg=act_cfg, drop_rate=drop_rate, memory_efficient=memory_efficient) self.stages.append(stage) channels += depth * self.growth_rate if i != self.num_stages - 1: transition = DenseTransition( in_channels=channels, out_channels=math.floor(channels * compression_factor), norm_cfg=norm_cfg, act_cfg=act_cfg, ) channels = math.floor(channels * compression_factor) else: # Final layers after dense block is just bn with act. # Unlike the paper, the original repo also put this in # transition layer, whereas torchvision take this out. # We reckon this as transition layer here. transition = nn.Sequential( build_norm_layer(norm_cfg, channels)[1], self.act, ) self.transitions.append(transition) self._freeze_stages()
def __init__(self, in_channels, out_channels, dim_after_concat, act_cfg=dict(type='ReLU'), upsample_cfg=dict(type='nearest', scale_factor=2), sn_eps=1e-6, sn_style='ajbrock', bn_eps=1e-5, with_spectral_norm=True, input_is_label=False, auto_sync_bn=True, channel_ratio=4): super().__init__() self.in_channels = in_channels self.out_channels = out_channels self.hidden_channels = self.in_channels // channel_ratio self.activation = build_activation_layer(act_cfg) self.upsample_cfg = deepcopy(upsample_cfg) self.with_upsample = upsample_cfg is not None if self.with_upsample: self.upsample_layer = build_upsample_layer(self.upsample_cfg) # Here in_channels of BigGANGenResBlock equal to num_features of # BigGANConditionBN self.bn1 = BigGANConditionBN(in_channels, dim_after_concat, sn_eps=sn_eps, sn_style=sn_style, bn_eps=bn_eps, input_is_label=input_is_label, with_spectral_norm=with_spectral_norm, auto_sync_bn=auto_sync_bn) # Here out_channels of BigGANGenResBlock equal to num_features of # BigGANConditionBN self.bn2 = BigGANConditionBN(self.hidden_channels, dim_after_concat, sn_eps=sn_eps, sn_style=sn_style, bn_eps=bn_eps, input_is_label=input_is_label, with_spectral_norm=with_spectral_norm, auto_sync_bn=auto_sync_bn) self.bn3 = BigGANConditionBN(self.hidden_channels, dim_after_concat, sn_eps=sn_eps, sn_style=sn_style, bn_eps=bn_eps, input_is_label=input_is_label, with_spectral_norm=with_spectral_norm, auto_sync_bn=auto_sync_bn) self.bn4 = BigGANConditionBN(self.hidden_channels, dim_after_concat, sn_eps=sn_eps, sn_style=sn_style, bn_eps=bn_eps, input_is_label=input_is_label, with_spectral_norm=with_spectral_norm, auto_sync_bn=auto_sync_bn) self.conv1 = SNConvModule(in_channels=in_channels, out_channels=self.hidden_channels, kernel_size=1, stride=1, padding=0, act_cfg=None, with_spectral_norm=with_spectral_norm, spectral_norm_cfg=dict(eps=sn_eps, sn_style=sn_style)) self.conv2 = SNConvModule(in_channels=self.hidden_channels, out_channels=self.hidden_channels, kernel_size=3, stride=1, padding=1, act_cfg=None, with_spectral_norm=with_spectral_norm, spectral_norm_cfg=dict(eps=sn_eps, sn_style=sn_style)) self.conv3 = SNConvModule(in_channels=self.hidden_channels, out_channels=self.hidden_channels, kernel_size=3, stride=1, padding=1, act_cfg=None, with_spectral_norm=with_spectral_norm, spectral_norm_cfg=dict(eps=sn_eps, sn_style=sn_style)) self.conv4 = SNConvModule(in_channels=self.hidden_channels, out_channels=out_channels, kernel_size=1, stride=1, padding=0, act_cfg=None, with_spectral_norm=with_spectral_norm, spectral_norm_cfg=dict(eps=sn_eps, sn_style=sn_style))
def __init__(self, arch='768/32', in_channels=3, norm_cfg=dict(type='BN'), act_cfg=dict(type='GELU'), out_indices=-1, frozen_stages=0, init_cfg=None): super().__init__(init_cfg=init_cfg) if isinstance(arch, str): assert arch in self.arch_settings, \ f'Unavailable arch, please choose from ' \ f'({set(self.arch_settings)}) or pass a dict.' arch = self.arch_settings[arch] elif isinstance(arch, dict): essential_keys = { 'embed_dims', 'depth', 'patch_size', 'kernel_size' } assert isinstance(arch, dict) and essential_keys <= set(arch), \ f'Custom arch needs a dict with keys {essential_keys}' self.embed_dims = arch['embed_dims'] self.depth = arch['depth'] self.patch_size = arch['patch_size'] self.kernel_size = arch['kernel_size'] self.act = build_activation_layer(act_cfg) # check out indices and frozen stages if isinstance(out_indices, int): out_indices = [out_indices] assert isinstance(out_indices, Sequence), \ f'"out_indices" must by a sequence or int, ' \ f'get {type(out_indices)} instead.' for i, index in enumerate(out_indices): if index < 0: out_indices[i] = self.depth + index assert out_indices[i] >= 0, f'Invalid out_indices {index}' self.out_indices = out_indices self.frozen_stages = frozen_stages # Set stem layers self.stem = nn.Sequential( nn.Conv2d(in_channels, self.embed_dims, kernel_size=self.patch_size, stride=self.patch_size), self.act, build_norm_layer(norm_cfg, self.embed_dims)[1]) # Set conv2d according to torch version convfunc = nn.Conv2d if digit_version(torch.__version__) < digit_version('1.9.0'): convfunc = Conv2dAdaptivePadding # Repetitions of ConvMixer Layer self.stages = nn.Sequential(*[ nn.Sequential( Residual( nn.Sequential( convfunc(self.embed_dims, self.embed_dims, self.kernel_size, groups=self.embed_dims, padding='same'), self.act, build_norm_layer(norm_cfg, self.embed_dims)[1])), nn.Conv2d(self.embed_dims, self.embed_dims, kernel_size=1), self.act, build_norm_layer(norm_cfg, self.embed_dims)[1]) for _ in range(self.depth) ]) self._freeze_stages()