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()
