def __init__(self,
in_channels: int,
expansion: float,
out_channels: int,
bias_attr=False):
super().__init__()
self.in_channels = in_channels
self.expansion = expansion
self.out_channels = out_channels
self.bottle_channels = round(self.expansion * self.in_channels)
self.body = nn.Sequential(
# pw
Conv2DNormLReLU(self.in_channels,
self.bottle_channels,
kernel_size=1,
bias_attr=bias_attr),
# dw
nn.Conv2D(self.bottle_channels,
self.bottle_channels,
kernel_size=3,
stride=1,
padding=0,
groups=self.bottle_channels,
bias_attr=True),
nn.GroupNorm(1, self.bottle_channels),
nn.LeakyReLU(0.2),
# pw & linear
nn.Conv2D(self.bottle_channels,
self.out_channels,
kernel_size=1,
padding=0,
bias_attr=False),
nn.GroupNorm(1, self.out_channels),
)
def __init__(self, channel: int = 64, nblocks: int = 3) -> None:
super().__init__()
channel = channel // 2
last_channel = channel
f = [
spectral_norm(
nn.Conv2D(3, channel, 3, stride=1, padding=1,
bias_attr=False)),
nn.LeakyReLU(0.2)
]
in_h = 256
for i in range(1, nblocks):
f.extend([
spectral_norm(
nn.Conv2D(last_channel,
channel * 2,
3,
stride=2,
padding=1,
bias_attr=False)),
nn.LeakyReLU(0.2),
spectral_norm(
nn.Conv2D(channel * 2,
channel * 4,
3,
stride=1,
padding=1,
bias_attr=False)),
nn.GroupNorm(1, channel * 4),
nn.LeakyReLU(0.2)
])
last_channel = channel * 4
channel = channel * 2
in_h = in_h // 2
self.body = nn.Sequential(*f)
self.head = nn.Sequential(*[
spectral_norm(
nn.Conv2D(last_channel,
channel * 2,
3,
stride=1,
padding=1,
bias_attr=False)),
nn.GroupNorm(1, channel * 2),
nn.LeakyReLU(0.2),
spectral_norm(
nn.Conv2D(
channel * 2, 1, 3, stride=1, padding=1, bias_attr=False))
])
def group_norm(out_channels):
"""group normal function
Args:
out_channels (int): out channel nums
Returns:
nn.Module: GroupNorm op
"""
num_groups = 32
if out_channels % 32 == 0:
return nn.GroupNorm(num_groups, out_channels)
else:
return nn.GroupNorm(num_groups // 2, out_channels)
def __init__(self,
in_channel,
out_channel,
time_dim,
use_affine_time=False,
dropout=0):
super().__init__()
self.use_affine_time = use_affine_time
time_out_dim = out_channel
time_scale = 1
norm_affine = None
if self.use_affine_time:
time_out_dim *= 2
time_scale = 1e-10
norm_affine = False
self.norm1 = nn.GroupNorm(32, in_channel)
self.activation1 = Swish()
self.conv1 = conv2d(in_channel, out_channel, 3, padding=1)
self.time = nn.Sequential(
Swish(), linear(time_dim, time_out_dim, scale=time_scale))
# A bug if we set `weight_attr` and `bias_attr` to False.
# Delete the weights to fix it as a temporary solution.
# self.norm2 = nn.GroupNorm(32, out_channel, weight_attr=norm_affine, bias_attr=norm_affine)
self.norm2 = nn.GroupNorm(32, out_channel)
if self.use_affine_time:
del self.norm2.weight
del self.norm2.bias
self.norm2.weight = None
self.norm2.bias = None
self.activation2 = Swish()
self.dropout = nn.Dropout(dropout)
self.conv2 = conv2d(out_channel,
out_channel,
3,
padding=1,
scale=1e-10)
if in_channel != out_channel:
self.skip = conv2d(in_channel, out_channel, 1)
else:
self.skip = None
def __init__(self, in_channel, n_head=1):
super().__init__()
self.n_head = n_head
self.norm = nn.GroupNorm(32, in_channel)
self.qkv = conv2d(in_channel, in_channel * 3, 1)
self.out = conv2d(in_channel, in_channel, 1, scale=1e-10)
self.key_dim = in_channel
def __init__(self,
ch_in,
ch_out,
filter_size,
stride=1,
groups=1,
norm_type=None,
norm_groups=32,
norm_decay=0.,
freeze_norm=False,
act=None):
super(ConvNormLayer, self).__init__()
self.act = act
norm_lr = 0. if freeze_norm else 1.
if norm_type is not None:
assert norm_type in ['bn', 'sync_bn', 'gn'], \
"norm_type should be one of ['bn', 'sync_bn', 'gn'], but got {}".format(norm_type)
param_attr = ParamAttr(
initializer=Constant(1.0),
learning_rate=norm_lr,
regularizer=L2Decay(norm_decay),
)
bias_attr = ParamAttr(learning_rate=norm_lr,
regularizer=L2Decay(norm_decay))
global_stats = True if freeze_norm else None
if norm_type in ['bn', 'sync_bn']:
self.norm = nn.BatchNorm2D(
ch_out,
weight_attr=param_attr,
bias_attr=bias_attr,
use_global_stats=global_stats,
)
elif norm_type == 'gn':
self.norm = nn.GroupNorm(num_groups=norm_groups,
num_channels=ch_out,
weight_attr=param_attr,
bias_attr=bias_attr)
norm_params = self.norm.parameters()
if freeze_norm:
for param in norm_params:
param.stop_gradient = True
conv_bias_attr = False
else:
conv_bias_attr = True
self.norm = None
self.conv = nn.Conv2D(
in_channels=ch_in,
out_channels=ch_out,
kernel_size=filter_size,
stride=stride,
padding=(filter_size - 1) // 2,
groups=groups,
weight_attr=ParamAttr(initializer=Normal(mean=0., std=0.001)),
bias_attr=conv_bias_attr)
def __init__(self,
ch_in,
ch_out,
filter_size,
stride=1,
norm_type='bn',
norm_groups=32,
use_dcn=False,
norm_decay=0.,
freeze_norm=False,
act=None,
name=None):
super(ConvNormLayer, self).__init__()
assert norm_type in ['bn', 'sync_bn', 'gn']
self.act = act
self.conv = nn.Conv2D(in_channels=ch_in,
out_channels=ch_out,
kernel_size=filter_size,
stride=stride,
padding=(filter_size - 1) // 2,
groups=1,
weight_attr=ParamAttr(name=name + "_weights",
initializer=Normal(
mean=0., std=0.01)),
bias_attr=False)
norm_lr = 0. if freeze_norm else 1.
norm_name = name + '_bn'
param_attr = ParamAttr(name=norm_name + "_scale",
learning_rate=norm_lr,
regularizer=L2Decay(norm_decay))
bias_attr = ParamAttr(name=norm_name + "_offset",
learning_rate=norm_lr,
regularizer=L2Decay(norm_decay))
global_stats = True if freeze_norm else False
if norm_type in ['bn', 'sync_bn']:
self.norm = nn.BatchNorm(ch_out,
param_attr=param_attr,
bias_attr=bias_attr,
use_global_stats=global_stats,
moving_mean_name=norm_name + '_mean',
moving_variance_name=norm_name +
'_variance')
elif norm_type == 'gn':
self.norm = nn.GroupNorm(num_groups=norm_groups,
num_channels=ch_out,
weight_attr=param_attr,
bias_attr=bias_attr)
norm_params = self.norm.parameters()
if freeze_norm:
for param in norm_params:
param.stop_gradient = True
def __init__(self,
ch_in,
ch_out,
filter_size,
stride,
norm_type='bn',
norm_groups=32,
use_dcn=False,
norm_name=None,
bias_on=False,
lr_scale=1.,
name=None):
super(ConvNormLayer, self).__init__()
assert norm_type in ['bn', 'sync_bn', 'gn']
if bias_on:
bias_attr = ParamAttr(name=name + "_bias",
initializer=Constant(value=0.),
learning_rate=lr_scale)
else:
bias_attr = False
self.conv = nn.Conv2D(in_channels=ch_in,
out_channels=ch_out,
kernel_size=filter_size,
stride=stride,
padding=(filter_size - 1) // 2,
groups=1,
weight_attr=ParamAttr(name=name + "_weight",
initializer=Normal(
mean=0., std=0.01),
learning_rate=1.),
bias_attr=bias_attr)
param_attr = ParamAttr(name=norm_name + "_scale",
learning_rate=1.,
regularizer=L2Decay(0.))
bias_attr = ParamAttr(name=norm_name + "_offset",
learning_rate=1.,
regularizer=L2Decay(0.))
if norm_type in ['bn', 'sync_bn']:
self.norm = nn.BatchNorm2D(ch_out,
weight_attr=param_attr,
bias_attr=bias_attr)
elif norm_type == 'gn':
self.norm = nn.GroupNorm(num_groups=norm_groups,
num_channels=ch_out,
weight_attr=param_attr,
bias_attr=bias_attr)
def _make_layers(self,
in_dims,
out_dims,
kernel_size,
num_groups,
weight_attr=None,
bias_attr=None):
return nn.Sequential(
nn.Conv2D(in_dims,
out_dims,
kernel_size,
padding=kernel_size // 2,
weight_attr=weight_attr,
bias_attr=bias_attr), nn.GroupNorm(num_groups, out_dims),
nn.ReLU())
def __init__(self, config, layer_id=0):
super().__init__()
self.in_conv_dim = config.conv_dim[layer_id] if layer_id > 0 else 1
self.out_conv_dim = config.conv_dim[layer_id]
self.conv = nn.Conv1D(
self.in_conv_dim,
self.out_conv_dim,
kernel_size=config.conv_kernel[layer_id],
stride=config.conv_stride[layer_id],
bias_attr=config.conv_bias,
)
self.activation = ACT2FN[config.feat_extract_activation]
self.layer_norm = nn.GroupNorm(num_groups=self.out_conv_dim,
num_channels=self.out_conv_dim)
def __init__(self,
in_channels: int,
out_channels: int,
kernel_size: int = 3,
stride: int = 1,
padding: int = 1,
bias_attr=False) -> None:
super().__init__()
self.conv = nn.Conv2D(in_channels,
out_channels,
kernel_size,
stride,
padding,
bias_attr=bias_attr)
# NOTE layer norm is crucial for animegan!
self.norm = nn.GroupNorm(1, out_channels)
self.lrelu = nn.LeakyReLU(0.2)
def __init__(self,
in_channels,
out_channels=None,
kernel_size=3,
norm_type='bn',
norm_groups=32,
act='swish'):
super(SeparableConvLayer, self).__init__()
assert norm_type in ['bn', 'sync_bn', 'gn', None]
assert act in ['swish', 'relu', None]
self.in_channels = in_channels
if out_channels is None:
self.out_channels = self.in_channels
self.norm_type = norm_type
self.norm_groups = norm_groups
self.depthwise_conv = nn.Conv2D(in_channels,
in_channels,
kernel_size,
padding=kernel_size // 2,
groups=in_channels,
bias_attr=False)
self.pointwise_conv = nn.Conv2D(in_channels, self.out_channels, 1)
# norm type
if self.norm_type == 'bn':
self.norm = nn.BatchNorm2D(self.out_channels)
elif self.norm_type == 'sync_bn':
self.norm = nn.SyncBatchNorm(self.out_channels)
elif self.norm_type == 'gn':
self.norm = nn.GroupNorm(num_groups=self.norm_groups,
num_channels=self.out_channels)
# activation
if act == 'swish':
self.act = nn.Swish()
elif act == 'relu':
self.act = nn.ReLU()
def __init__(self, channel, groups=64):
super(sa_layer, self).__init__()
self.groups = groups
self.avg_pool = nn.AdaptiveAvgPool2D(1)
self.cweight = self.create_parameter(
shape=[1, channel // (2 * groups), 1, 1],
default_initializer=paddle.nn.initializer.Assign(
paddle.zeros([1, channel // (2 * groups), 1, 1])))
self.cbias = self.create_parameter(
shape=[1, channel // (2 * groups), 1, 1],
default_initializer=paddle.nn.initializer.Assign(
paddle.ones([1, channel // (2 * groups), 1, 1])))
self.sweight = self.create_parameter(
shape=[1, channel // (2 * groups), 1, 1],
default_initializer=paddle.nn.initializer.Assign(
paddle.zeros([1, channel // (2 * groups), 1, 1])))
self.sbias = self.create_parameter(
shape=[1, channel // (2 * groups), 1, 1],
default_initializer=paddle.nn.initializer.Assign(
paddle.ones([1, channel // (2 * groups), 1, 1])))
self.sigmoid = nn.Sigmoid()
self.gn = nn.GroupNorm(channel // (2 * groups),
channel // (2 * groups))
def __init__(self,
num_queries=300,
position_embed_type='sine',
return_intermediate_dec=True,
backbone_num_channels=[512, 1024, 2048],
num_feature_levels=4,
num_encoder_points=4,
num_decoder_points=4,
hidden_dim=256,
nhead=8,
num_encoder_layers=6,
num_decoder_layers=6,
dim_feedforward=1024,
dropout=0.1,
activation="relu",
lr_mult=0.1,
weight_attr=None,
bias_attr=None):
super(DeformableTransformer, self).__init__()
assert position_embed_type in ['sine', 'learned'], \
f'ValueError: position_embed_type not supported {position_embed_type}!'
assert len(backbone_num_channels) <= num_feature_levels
self.hidden_dim = hidden_dim
self.nhead = nhead
self.num_feature_levels = num_feature_levels
encoder_layer = DeformableTransformerEncoderLayer(
hidden_dim, nhead, dim_feedforward, dropout, activation,
num_feature_levels, num_encoder_points, weight_attr, bias_attr)
self.encoder = DeformableTransformerEncoder(encoder_layer,
num_encoder_layers)
decoder_layer = DeformableTransformerDecoderLayer(
hidden_dim, nhead, dim_feedforward, dropout, activation,
num_feature_levels, num_decoder_points, weight_attr, bias_attr)
self.decoder = DeformableTransformerDecoder(decoder_layer,
num_decoder_layers,
return_intermediate_dec)
self.level_embed = nn.Embedding(num_feature_levels, hidden_dim)
self.tgt_embed = nn.Embedding(num_queries, hidden_dim)
self.query_pos_embed = nn.Embedding(num_queries, hidden_dim)
self.reference_points = nn.Linear(
hidden_dim,
2,
weight_attr=ParamAttr(learning_rate=lr_mult),
bias_attr=ParamAttr(learning_rate=lr_mult))
self.input_proj = nn.LayerList()
for in_channels in backbone_num_channels:
self.input_proj.append(
nn.Sequential(
nn.Conv2D(in_channels,
hidden_dim,
kernel_size=1,
weight_attr=weight_attr,
bias_attr=bias_attr),
nn.GroupNorm(32, hidden_dim)))
in_channels = backbone_num_channels[-1]
for _ in range(num_feature_levels - len(backbone_num_channels)):
self.input_proj.append(
nn.Sequential(
nn.Conv2D(in_channels,
hidden_dim,
kernel_size=3,
stride=2,
padding=1,
weight_attr=weight_attr,
bias_attr=bias_attr),
nn.GroupNorm(32, hidden_dim)))
in_channels = hidden_dim
self.position_embedding = PositionEmbedding(
hidden_dim // 2,
normalize=True if position_embed_type == 'sine' else False,
embed_type=position_embed_type,
offset=-0.5)
self._reset_parameters()
def func_test_layer_str(self):
module = nn.ELU(0.2)
self.assertEqual(str(module), 'ELU(alpha=0.2)')
module = nn.CELU(0.2)
self.assertEqual(str(module), 'CELU(alpha=0.2)')
module = nn.GELU(True)
self.assertEqual(str(module), 'GELU(approximate=True)')
module = nn.Hardshrink()
self.assertEqual(str(module), 'Hardshrink(threshold=0.5)')
module = nn.Hardswish(name="Hardswish")
self.assertEqual(str(module), 'Hardswish(name=Hardswish)')
module = nn.Tanh(name="Tanh")
self.assertEqual(str(module), 'Tanh(name=Tanh)')
module = nn.Hardtanh(name="Hardtanh")
self.assertEqual(str(module),
'Hardtanh(min=-1.0, max=1.0, name=Hardtanh)')
module = nn.PReLU(1, 0.25, name="PReLU", data_format="NCHW")
self.assertEqual(
str(module),
'PReLU(num_parameters=1, data_format=NCHW, init=0.25, dtype=float32, name=PReLU)'
)
module = nn.ReLU()
self.assertEqual(str(module), 'ReLU()')
module = nn.ReLU6()
self.assertEqual(str(module), 'ReLU6()')
module = nn.SELU()
self.assertEqual(
str(module),
'SELU(scale=1.0507009873554805, alpha=1.6732632423543772)')
module = nn.LeakyReLU()
self.assertEqual(str(module), 'LeakyReLU(negative_slope=0.01)')
module = nn.Sigmoid()
self.assertEqual(str(module), 'Sigmoid()')
module = nn.Hardsigmoid()
self.assertEqual(str(module), 'Hardsigmoid()')
module = nn.Softplus()
self.assertEqual(str(module), 'Softplus(beta=1, threshold=20)')
module = nn.Softshrink()
self.assertEqual(str(module), 'Softshrink(threshold=0.5)')
module = nn.Softsign()
self.assertEqual(str(module), 'Softsign()')
module = nn.Swish()
self.assertEqual(str(module), 'Swish()')
module = nn.Tanhshrink()
self.assertEqual(str(module), 'Tanhshrink()')
module = nn.ThresholdedReLU()
self.assertEqual(str(module), 'ThresholdedReLU(threshold=1.0)')
module = nn.LogSigmoid()
self.assertEqual(str(module), 'LogSigmoid()')
module = nn.Softmax()
self.assertEqual(str(module), 'Softmax(axis=-1)')
module = nn.LogSoftmax()
self.assertEqual(str(module), 'LogSoftmax(axis=-1)')
module = nn.Maxout(groups=2)
self.assertEqual(str(module), 'Maxout(groups=2, axis=1)')
module = nn.Linear(2, 4, name='linear')
self.assertEqual(
str(module),
'Linear(in_features=2, out_features=4, dtype=float32, name=linear)'
)
module = nn.Upsample(size=[12, 12])
self.assertEqual(
str(module),
'Upsample(size=[12, 12], mode=nearest, align_corners=False, align_mode=0, data_format=NCHW)'
)
module = nn.UpsamplingNearest2D(size=[12, 12])
self.assertEqual(
str(module),
'UpsamplingNearest2D(size=[12, 12], data_format=NCHW)')
module = nn.UpsamplingBilinear2D(size=[12, 12])
self.assertEqual(
str(module),
'UpsamplingBilinear2D(size=[12, 12], data_format=NCHW)')
module = nn.Bilinear(in1_features=5, in2_features=4, out_features=1000)
self.assertEqual(
str(module),
'Bilinear(in1_features=5, in2_features=4, out_features=1000, dtype=float32)'
)
module = nn.Dropout(p=0.5)
self.assertEqual(str(module),
'Dropout(p=0.5, axis=None, mode=upscale_in_train)')
module = nn.Dropout2D(p=0.5)
self.assertEqual(str(module), 'Dropout2D(p=0.5, data_format=NCHW)')
module = nn.Dropout3D(p=0.5)
self.assertEqual(str(module), 'Dropout3D(p=0.5, data_format=NCDHW)')
module = nn.AlphaDropout(p=0.5)
self.assertEqual(str(module), 'AlphaDropout(p=0.5)')
module = nn.Pad1D(padding=[1, 2], mode='constant')
self.assertEqual(
str(module),
'Pad1D(padding=[1, 2], mode=constant, value=0.0, data_format=NCL)')
module = nn.Pad2D(padding=[1, 0, 1, 2], mode='constant')
self.assertEqual(
str(module),
'Pad2D(padding=[1, 0, 1, 2], mode=constant, value=0.0, data_format=NCHW)'
)
module = nn.ZeroPad2D(padding=[1, 0, 1, 2])
self.assertEqual(str(module),
'ZeroPad2D(padding=[1, 0, 1, 2], data_format=NCHW)')
module = nn.Pad3D(padding=[1, 0, 1, 2, 0, 0], mode='constant')
self.assertEqual(
str(module),
'Pad3D(padding=[1, 0, 1, 2, 0, 0], mode=constant, value=0.0, data_format=NCDHW)'
)
module = nn.CosineSimilarity(axis=0)
self.assertEqual(str(module), 'CosineSimilarity(axis=0, eps=1e-08)')
module = nn.Embedding(10, 3, sparse=True)
self.assertEqual(str(module), 'Embedding(10, 3, sparse=True)')
module = nn.Conv1D(3, 2, 3)
self.assertEqual(str(module),
'Conv1D(3, 2, kernel_size=[3], data_format=NCL)')
module = nn.Conv1DTranspose(2, 1, 2)
self.assertEqual(
str(module),
'Conv1DTranspose(2, 1, kernel_size=[2], data_format=NCL)')
module = nn.Conv2D(4, 6, (3, 3))
self.assertEqual(str(module),
'Conv2D(4, 6, kernel_size=[3, 3], data_format=NCHW)')
module = nn.Conv2DTranspose(4, 6, (3, 3))
self.assertEqual(
str(module),
'Conv2DTranspose(4, 6, kernel_size=[3, 3], data_format=NCHW)')
module = nn.Conv3D(4, 6, (3, 3, 3))
self.assertEqual(
str(module),
'Conv3D(4, 6, kernel_size=[3, 3, 3], data_format=NCDHW)')
module = nn.Conv3DTranspose(4, 6, (3, 3, 3))
self.assertEqual(
str(module),
'Conv3DTranspose(4, 6, kernel_size=[3, 3, 3], data_format=NCDHW)')
module = nn.PairwiseDistance()
self.assertEqual(str(module), 'PairwiseDistance(p=2.0)')
module = nn.InstanceNorm1D(2)
self.assertEqual(str(module),
'InstanceNorm1D(num_features=2, epsilon=1e-05)')
module = nn.InstanceNorm2D(2)
self.assertEqual(str(module),
'InstanceNorm2D(num_features=2, epsilon=1e-05)')
module = nn.InstanceNorm3D(2)
self.assertEqual(str(module),
'InstanceNorm3D(num_features=2, epsilon=1e-05)')
module = nn.GroupNorm(num_channels=6, num_groups=6)
self.assertEqual(
str(module),
'GroupNorm(num_groups=6, num_channels=6, epsilon=1e-05)')
module = nn.LayerNorm([2, 2, 3])
self.assertEqual(
str(module),
'LayerNorm(normalized_shape=[2, 2, 3], epsilon=1e-05)')
module = nn.BatchNorm1D(1)
self.assertEqual(
str(module),
'BatchNorm1D(num_features=1, momentum=0.9, epsilon=1e-05, data_format=NCL)'
)
module = nn.BatchNorm2D(1)
self.assertEqual(
str(module),
'BatchNorm2D(num_features=1, momentum=0.9, epsilon=1e-05)')
module = nn.BatchNorm3D(1)
self.assertEqual(
str(module),
'BatchNorm3D(num_features=1, momentum=0.9, epsilon=1e-05, data_format=NCDHW)'
)
module = nn.SyncBatchNorm(2)
self.assertEqual(
str(module),
'SyncBatchNorm(num_features=2, momentum=0.9, epsilon=1e-05)')
module = nn.LocalResponseNorm(size=5)
self.assertEqual(
str(module),
'LocalResponseNorm(size=5, alpha=0.0001, beta=0.75, k=1.0)')
module = nn.AvgPool1D(kernel_size=2, stride=2, padding=0)
self.assertEqual(str(module),
'AvgPool1D(kernel_size=2, stride=2, padding=0)')
module = nn.AvgPool2D(kernel_size=2, stride=2, padding=0)
self.assertEqual(str(module),
'AvgPool2D(kernel_size=2, stride=2, padding=0)')
module = nn.AvgPool3D(kernel_size=2, stride=2, padding=0)
self.assertEqual(str(module),
'AvgPool3D(kernel_size=2, stride=2, padding=0)')
module = nn.MaxPool1D(kernel_size=2, stride=2, padding=0)
self.assertEqual(str(module),
'MaxPool1D(kernel_size=2, stride=2, padding=0)')
module = nn.MaxPool2D(kernel_size=2, stride=2, padding=0)
self.assertEqual(str(module),
'MaxPool2D(kernel_size=2, stride=2, padding=0)')
module = nn.MaxPool3D(kernel_size=2, stride=2, padding=0)
self.assertEqual(str(module),
'MaxPool3D(kernel_size=2, stride=2, padding=0)')
module = nn.AdaptiveAvgPool1D(output_size=16)
self.assertEqual(str(module), 'AdaptiveAvgPool1D(output_size=16)')
module = nn.AdaptiveAvgPool2D(output_size=3)
self.assertEqual(str(module), 'AdaptiveAvgPool2D(output_size=3)')
module = nn.AdaptiveAvgPool3D(output_size=3)
self.assertEqual(str(module), 'AdaptiveAvgPool3D(output_size=3)')
module = nn.AdaptiveMaxPool1D(output_size=16, return_mask=True)
self.assertEqual(
str(module), 'AdaptiveMaxPool1D(output_size=16, return_mask=True)')
module = nn.AdaptiveMaxPool2D(output_size=3, return_mask=True)
self.assertEqual(str(module),
'AdaptiveMaxPool2D(output_size=3, return_mask=True)')
module = nn.AdaptiveMaxPool3D(output_size=3, return_mask=True)
self.assertEqual(str(module),
'AdaptiveMaxPool3D(output_size=3, return_mask=True)')
module = nn.SimpleRNNCell(16, 32)
self.assertEqual(str(module), 'SimpleRNNCell(16, 32)')
module = nn.LSTMCell(16, 32)
self.assertEqual(str(module), 'LSTMCell(16, 32)')
module = nn.GRUCell(16, 32)
self.assertEqual(str(module), 'GRUCell(16, 32)')
module = nn.PixelShuffle(3)
self.assertEqual(str(module), 'PixelShuffle(upscale_factor=3)')
module = nn.SimpleRNN(16, 32, 2)
self.assertEqual(
str(module),
'SimpleRNN(16, 32, num_layers=2\n (0): RNN(\n (cell): SimpleRNNCell(16, 32)\n )\n (1): RNN(\n (cell): SimpleRNNCell(32, 32)\n )\n)'
)
module = nn.LSTM(16, 32, 2)
self.assertEqual(
str(module),
'LSTM(16, 32, num_layers=2\n (0): RNN(\n (cell): LSTMCell(16, 32)\n )\n (1): RNN(\n (cell): LSTMCell(32, 32)\n )\n)'
)
module = nn.GRU(16, 32, 2)
self.assertEqual(
str(module),
'GRU(16, 32, num_layers=2\n (0): RNN(\n (cell): GRUCell(16, 32)\n )\n (1): RNN(\n (cell): GRUCell(32, 32)\n )\n)'
)
module1 = nn.Sequential(
('conv1', nn.Conv2D(1, 20, 5)), ('relu1', nn.ReLU()),
('conv2', nn.Conv2D(20, 64, 5)), ('relu2', nn.ReLU()))
self.assertEqual(
str(module1),
'Sequential(\n '\
'(conv1): Conv2D(1, 20, kernel_size=[5, 5], data_format=NCHW)\n '\
'(relu1): ReLU()\n '\
'(conv2): Conv2D(20, 64, kernel_size=[5, 5], data_format=NCHW)\n '\
'(relu2): ReLU()\n)'
)
module2 = nn.Sequential(
nn.Conv3DTranspose(4, 6, (3, 3, 3)),
nn.AvgPool3D(kernel_size=2, stride=2, padding=0),
nn.Tanh(name="Tanh"), module1, nn.Conv3D(4, 6, (3, 3, 3)),
nn.MaxPool3D(kernel_size=2, stride=2, padding=0), nn.GELU(True))
self.assertEqual(
str(module2),
'Sequential(\n '\
'(0): Conv3DTranspose(4, 6, kernel_size=[3, 3, 3], data_format=NCDHW)\n '\
'(1): AvgPool3D(kernel_size=2, stride=2, padding=0)\n '\
'(2): Tanh(name=Tanh)\n '\
'(3): Sequential(\n (conv1): Conv2D(1, 20, kernel_size=[5, 5], data_format=NCHW)\n (relu1): ReLU()\n'\
' (conv2): Conv2D(20, 64, kernel_size=[5, 5], data_format=NCHW)\n (relu2): ReLU()\n )\n '\
'(4): Conv3D(4, 6, kernel_size=[3, 3, 3], data_format=NCDHW)\n '\
'(5): MaxPool3D(kernel_size=2, stride=2, padding=0)\n '\
'(6): GELU(approximate=True)\n)'
)
def __init__(self,
ch_in,
ch_out,
filter_size,
stride,
groups=1,
norm_type='bn',
norm_decay=0.,
norm_groups=32,
use_dcn=False,
bias_on=False,
lr_scale=1.,
freeze_norm=False,
initializer=Normal(mean=0., std=0.01)):
super(ConvNormLayer, self).__init__()
assert norm_type in ['bn', 'sync_bn', 'gn']
if bias_on:
bias_attr = ParamAttr(initializer=Constant(value=0.),
learning_rate=lr_scale)
else:
bias_attr = False
if not use_dcn:
self.conv = nn.Conv2D(in_channels=ch_in,
out_channels=ch_out,
kernel_size=filter_size,
stride=stride,
padding=(filter_size - 1) // 2,
groups=groups,
weight_attr=ParamAttr(
initializer=initializer,
learning_rate=1.),
bias_attr=bias_attr)
else:
# in FCOS-DCN head, specifically need learning_rate and regularizer
self.conv = DeformableConvV2(in_channels=ch_in,
out_channels=ch_out,
kernel_size=filter_size,
stride=stride,
padding=(filter_size - 1) // 2,
groups=groups,
weight_attr=ParamAttr(
initializer=initializer,
learning_rate=1.),
bias_attr=True,
lr_scale=2.,
regularizer=L2Decay(norm_decay))
norm_lr = 0. if freeze_norm else 1.
param_attr = ParamAttr(learning_rate=norm_lr,
regularizer=L2Decay(norm_decay))
bias_attr = ParamAttr(learning_rate=norm_lr,
regularizer=L2Decay(norm_decay))
if norm_type == 'bn':
self.norm = nn.BatchNorm2D(ch_out,
weight_attr=param_attr,
bias_attr=bias_attr)
elif norm_type == 'sync_bn':
self.norm = nn.SyncBatchNorm(ch_out,
weight_attr=param_attr,
bias_attr=bias_attr)
elif norm_type == 'gn':
self.norm = nn.GroupNorm(num_groups=norm_groups,
num_channels=ch_out,
weight_attr=param_attr,
bias_attr=bias_attr)
def __init__(
self,
in_channel,
channel,
channel_multiplier,
n_res_blocks,
attn_type='vanilla',
attn_strides=[],
attn_heads=1,
use_affine_time=False,
dropout=0,
fold=1,
):
super().__init__()
self.fold = fold
time_dim = channel * 4
n_block = len(channel_multiplier)
self.time = nn.Sequential(
TimeEmbedding(channel),
linear(channel, time_dim),
Swish(),
linear(time_dim, time_dim),
)
down_layers = [conv2d(in_channel * (fold**2), channel, 3, padding=1)]
feat_channels = [channel]
in_channel = channel
for i in range(n_block):
for _ in range(n_res_blocks):
channel_mult = channel * channel_multiplier[i]
down_layers.append(
ResBlockWithAttention(
in_channel,
channel_mult,
time_dim,
dropout,
use_attention=attn_type != 'none'
and 2**i in attn_strides,
attention_type=attn_type,
attention_head=attn_heads,
use_affine_time=use_affine_time,
))
feat_channels.append(channel_mult)
in_channel = channel_mult
if i != n_block - 1:
down_layers.append(Downsample(in_channel))
feat_channels.append(in_channel)
self.down = nn.LayerList(down_layers)
self.mid = nn.LayerList([
ResBlockWithAttention(
in_channel,
in_channel,
time_dim,
dropout=dropout,
use_attention=attn_type != 'none',
attention_type=attn_type,
attention_head=attn_heads,
use_affine_time=use_affine_time,
),
ResBlockWithAttention(
in_channel,
in_channel,
time_dim,
dropout=dropout,
use_affine_time=use_affine_time,
),
])
up_layers = []
for i in reversed(range(n_block)):
for _ in range(n_res_blocks + 1):
channel_mult = channel * channel_multiplier[i]
up_layers.append(
ResBlockWithAttention(
in_channel + feat_channels.pop(),
channel_mult,
time_dim,
dropout=dropout,
use_attention=attn_type != 'none'
and 2**i in attn_strides,
attention_type=attn_type,
attention_head=attn_heads,
use_affine_time=use_affine_time,
))
in_channel = channel_mult
if i != 0:
up_layers.append(Upsample(in_channel))
self.up = nn.LayerList(up_layers)
self.out = nn.Sequential(
nn.GroupNorm(32, in_channel),
Swish(),
conv2d(in_channel, 3 * (fold**2), 3, padding=1, scale=1e-10),
)
