def __init__(self) -> None:
super().__init__()
self.A = nn.Sequential(Conv2DNormLReLU(3, 32, 7, padding=3),
Conv2DNormLReLU(32, 64, stride=2),
Conv2DNormLReLU(64, 64))
self.B = nn.Sequential(Conv2DNormLReLU(64, 128, stride=2),
Conv2DNormLReLU(128, 128),
Conv2DNormLReLU(128, 128))
self.C = nn.Sequential(InvertedresBlock(128, 2, 256),
InvertedresBlock(256, 2, 256),
InvertedresBlock(256, 2, 256),
InvertedresBlock(256, 2, 256),
Conv2DNormLReLU(256, 128))
self.D = nn.Sequential(nn.Upsample(scale_factor=2, mode='bilinear'),
Conv2DNormLReLU(128, 128),
Conv2DNormLReLU(128, 128))
self.E = nn.Sequential(nn.Upsample(scale_factor=2, mode='bilinear'),
Conv2DNormLReLU(128, 64),
Conv2DNormLReLU(64, 64),
Conv2DNormLReLU(64, 32, 7, padding=3))
self.out = nn.Sequential(nn.Conv2D(32, 3, 1, bias_attr=False),
nn.Tanh())
def __init__(self, latent_dim, output_nc, size=64, ngf=64):
"""Construct a Deep Convolutional generator
Args:
latent_dim (int) -- the number of latent dimension
output_nc (int) -- the number of channels in output images
size (int) -- size of output tensor
ngf (int) -- the number of filters in the last conv layer
Refer to https://arxiv.org/abs/1511.06434
"""
super(DeepConvGenerator, self).__init__()
self.latent_dim = latent_dim
self.ngf = ngf
self.init_size = size // 4
self.l1 = nn.Sequential(
nn.Linear(latent_dim, ngf * 2 * self.init_size**2))
self.conv_blocks = nn.Sequential(
nn.BatchNorm2D(ngf * 2),
nn.Upsample(scale_factor=2),
nn.Conv2D(ngf * 2, ngf * 2, 3, stride=1, padding=1),
nn.BatchNorm2D(ngf * 2, 0.2),
nn.LeakyReLU(0.2),
nn.Upsample(scale_factor=2),
nn.Conv2D(ngf * 2, ngf, 3, stride=1, padding=1),
nn.BatchNorm2D(ngf, 0.2),
nn.LeakyReLU(0.2),
nn.Conv2D(ngf, output_nc, 3, stride=1, padding=1),
nn.Tanh(),
)
def __init__(self,
num_classes,
backbone_indices,
in_channels,
align_corners=False,
lr_multiple=10):
super().__init__()
inter_channels = 256
self.cls_head_norm = nn.LayerNorm(normalized_shape=in_channels,
epsilon=1e-6)
self.cls_head = nn.Sequential(
layers.ConvBNReLU(in_channels=in_channels,
out_channels=inter_channels,
kernel_size=3,
padding=1),
nn.Upsample(scale_factor=2, mode='bilinear'),
layers.ConvBNReLU(in_channels=inter_channels,
out_channels=inter_channels,
kernel_size=3,
padding=1),
nn.Upsample(scale_factor=2, mode='bilinear'),
layers.ConvBNReLU(in_channels=inter_channels,
out_channels=inter_channels,
kernel_size=3,
padding=1),
nn.Upsample(scale_factor=2, mode='bilinear'),
layers.ConvBNReLU(in_channels=inter_channels,
out_channels=inter_channels,
kernel_size=3,
padding=1),
nn.Conv2D(in_channels=inter_channels,
out_channels=num_classes,
kernel_size=1))
aux_head_nums = len(backbone_indices)
self.aux_head_norms = nn.LayerList(
[nn.LayerNorm(normalized_shape=in_channels, epsilon=1e-6)] *
aux_head_nums)
self.aux_heads = nn.LayerList([
nn.Sequential(
layers.ConvBNReLU(in_channels=in_channels,
out_channels=inter_channels,
kernel_size=3,
padding=1),
nn.Upsample(scale_factor=4, mode='bilinear'),
nn.Conv2D(in_channels=inter_channels,
out_channels=num_classes,
kernel_size=1))
] * aux_head_nums)
self.in_channels = in_channels
self.lr_multiple = lr_multiple
self.backbone_indices = backbone_indices
self.init_weight()
def __init__(self,
group_count=4,
n_init=128,
n_hid=256,
n_blk_per_group=2,
output_channels=3,
vocab_size=8192):
super(Decoder, self).__init__()
self.vocab_size = vocab_size
blk_range = range(n_blk_per_group)
n_layers = group_count * n_blk_per_group
self.blocks = nn.Sequential(
('input', nn.Conv2D(vocab_size, n_init, 1)),
('group_1',
nn.Sequential(
*[(f'block_{i + 1}',
DecoderBlock(n_init if i == 0 else 8 * n_hid,
8 * n_hid,
n_layers=n_layers)) for i in blk_range],
('upsample', nn.Upsample(scale_factor=2, mode='nearest')),
)),
('group_2',
nn.Sequential(
*[(f'block_{i + 1}',
DecoderBlock(8 * n_hid if i == 0 else 4 * n_hid,
4 * n_hid,
n_layers=n_layers)) for i in blk_range],
('upsample', nn.Upsample(scale_factor=2, mode='nearest')),
)),
('group_3',
nn.Sequential(
*[(f'block_{i + 1}',
DecoderBlock(4 * n_hid if i == 0 else 2 * n_hid,
2 * n_hid,
n_layers=n_layers)) for i in blk_range],
('upsample', nn.Upsample(scale_factor=2, mode='nearest')),
)),
('group_4',
nn.Sequential(
*[(f'block_{i + 1}',
DecoderBlock(2 * n_hid if i == 0 else 1 * n_hid,
1 * n_hid,
n_layers=n_layers)) for i in blk_range], )),
('output',
nn.Sequential(
('relu', nn.ReLU()),
('conv', nn.Conv2D(1 * n_hid, 2 * output_channels, 1)),
)),
)
def upsample(in_tens, out_HW=(64, 64)):
in_H, in_W = in_tens.shape[2], in_tens.shape[3]
scale_factor_H, scale_factor_W = 1. * out_HW[0] / in_H, 1. * out_HW[1] / in_W
return nn.Upsample(scale_factor=(scale_factor_H, scale_factor_W),
mode='bilinear',
align_corners=False)(in_tens)
def __init__(self, channel):
layers = [
nn.Upsample(scale_factor=2, mode="nearest"),
conv2d(channel, channel, 3, padding=1),
]
super().__init__(*layers)
def __init__(self, in_channels, bilinear=False):
super(UpSample, self).__init__()
if bilinear:
self.up = nn.Upsample(scale_factor=2,
mode='bilinear',
align_corners=True)
else:
self.up = nn.Conv2DTranspose(in_channels, in_channels, kernel_size=2, stride=2)
def __init__(self, in_channels: int, out_channels: int, kernel_size: int, stride: int, upsample=None):
super(UpsampleConvLayer, self).__init__()
self.upsample = upsample
if upsample:
self.upsample_layer = nn.Upsample(scale_factor=upsample)
self.pad = int(np.floor(kernel_size / 2))
if self.pad != 0:
self.reflection_pad = nn.Pad2D([self.pad, self.pad, self.pad, self.pad], mode='reflect')
self.conv2d = nn.Conv2D(in_channels, out_channels, kernel_size, stride)
def __init__(self, in_nf=3, nf=64, HR_in=False):
super(PredeblurResNetPyramid, self).__init__()
self.in_nf = in_nf
self.nf = nf
self.HR_in = True if HR_in else False
self.Leaky_relu = nn.LeakyReLU(negative_slope=0.1)
if self.HR_in:
self.conv_first_1 = nn.Conv2D(in_channels=self.in_nf,
out_channels=self.nf,
kernel_size=3,
stride=1,
padding=1)
self.conv_first_2 = nn.Conv2D(in_channels=self.nf,
out_channels=self.nf,
kernel_size=3,
stride=2,
padding=1)
self.conv_first_3 = nn.Conv2D(in_channels=self.nf,
out_channels=self.nf,
kernel_size=3,
stride=2,
padding=1)
else:
self.conv_first = nn.Conv2D(in_channels=self.in_nf,
out_channels=self.nf,
kernel_size=3,
stride=1,
padding=1)
self.RB_L1_1 = ResidualBlockNoBN(nf=self.nf)
self.RB_L1_2 = ResidualBlockNoBN(nf=self.nf)
self.RB_L1_3 = ResidualBlockNoBN(nf=self.nf)
self.RB_L1_4 = ResidualBlockNoBN(nf=self.nf)
self.RB_L1_5 = ResidualBlockNoBN(nf=self.nf)
self.RB_L2_1 = ResidualBlockNoBN(nf=self.nf)
self.RB_L2_2 = ResidualBlockNoBN(nf=self.nf)
self.RB_L3_1 = ResidualBlockNoBN(nf=self.nf)
self.deblur_L2_conv = nn.Conv2D(in_channels=self.nf,
out_channels=self.nf,
kernel_size=3,
stride=2,
padding=1)
self.deblur_L3_conv = nn.Conv2D(in_channels=self.nf,
out_channels=self.nf,
kernel_size=3,
stride=2,
padding=1)
self.upsample = nn.Upsample(scale_factor=2,
mode="bilinear",
align_corners=False,
align_mode=0)
def __init__(self,
depth_mult=1.0,
in_channels=[256, 512, 1024],
depthwise=False,
act='silu'):
super(YOLOCSPPAN, self).__init__()
self.in_channels = in_channels
self._out_channels = in_channels
Conv = DWConv if depthwise else BaseConv
self.upsample = nn.Upsample(scale_factor=2, mode="nearest")
# top-down fpn
self.lateral_convs = nn.LayerList()
self.fpn_blocks = nn.LayerList()
for idx in range(len(in_channels) - 1, 0, -1):
self.lateral_convs.append(
BaseConv(int(in_channels[idx]),
int(in_channels[idx - 1]),
1,
1,
act=act))
self.fpn_blocks.append(
CSPLayer(int(in_channels[idx - 1] * 2),
int(in_channels[idx - 1]),
round(3 * depth_mult),
shortcut=False,
depthwise=depthwise,
act=act))
# bottom-up pan
self.downsample_convs = nn.LayerList()
self.pan_blocks = nn.LayerList()
for idx in range(len(in_channels) - 1):
self.downsample_convs.append(
Conv(int(in_channels[idx]),
int(in_channels[idx]),
3,
stride=2,
act=act))
self.pan_blocks.append(
CSPLayer(int(in_channels[idx] * 2),
int(in_channels[idx + 1]),
round(3 * depth_mult),
shortcut=False,
depthwise=depthwise,
act=act))
def __init__(self,
num_classes,
backbone_indices,
in_channels,
mla_channels=256,
mlahead_channels=128,
lr_multiple=10):
super().__init__()
if len(backbone_indices) != 4:
raise RuntimeError
self.mla_feat_nums = len(backbone_indices)
self.norms = nn.LayerList(
[nn.LayerNorm(normalized_shape=in_channels, epsilon=1e-6)] *
self.mla_feat_nums)
self.mla = ConvMLA(in_channels, mla_channels)
self.aux_heads = nn.LayerList([
nn.Conv2D(in_channels=mla_channels,
out_channels=num_classes,
kernel_size=1)
] * self.mla_feat_nums)
self.feat_convs = nn.LayerList([
nn.Sequential(
layers.ConvBNReLU(in_channels=mla_channels,
out_channels=mlahead_channels,
kernel_size=3,
padding=1),
layers.ConvBNReLU(in_channels=mlahead_channels,
out_channels=mlahead_channels,
kernel_size=3,
padding=1),
nn.Upsample(
scale_factor=4, mode='bilinear', align_corners=True))
] * self.mla_feat_nums)
self.backbone_indices = backbone_indices
self.in_channels = in_channels
self.cls_head = nn.Conv2D(in_channels=4 * mlahead_channels,
out_channels=num_classes,
kernel_size=3,
padding=1)
def __init__(self,
in_channels,
out_channels,
n_cat,
use_deconv=False,
align_corners=False):
super(UpSampling, self).__init__()
if use_deconv:
self.up = nn.Conv2DTranspose(
in_channels, out_channels, kernel_size=2, stride=2, padding=0)
else:
self.up = nn.Sequential(
nn.Upsample(
scale_factor=2,
mode='bilinear',
align_corners=align_corners),
nn.Conv2D(in_channels, out_channels, 1, 1, 0))
self.conv = DoubleConv(n_cat * out_channels, out_channels)
def __init__(self, path=None, features=256, non_negative=True):
"""Init.
Args:
path (str, optional): Path to saved model. Defaults to None.
features (int, optional): Number of features. Defaults to 256.
backbone (str, optional): Backbone network for encoder. Defaults to resnet50
"""
print("Loading weights: ", path)
super(MidasNet, self).__init__()
use_pretrained = False if path is None else True
self.pretrained, self.scratch = _make_encoder(
backbone="resnext101_wsl",
features=features,
use_pretrained=use_pretrained)
self.scratch.refinenet4 = FeatureFusionBlock(features)
self.scratch.refinenet3 = FeatureFusionBlock(features)
self.scratch.refinenet2 = FeatureFusionBlock(features)
self.scratch.refinenet1 = FeatureFusionBlock(features)
output_conv = [
nn.Conv2D(features, 128, kernel_size=3, stride=1, padding=1),
nn.Upsample(scale_factor=2, mode="bilinear"),
nn.Conv2D(128, 32, kernel_size=3, stride=1, padding=1),
nn.ReLU(),
nn.Conv2D(32, 1, kernel_size=1, stride=1, padding=0),
nn.ReLU() if non_negative else nn.Identity(),
]
if non_negative:
output_conv.append(nn.ReLU())
self.scratch.output_conv = nn.Sequential(*output_conv)
if path:
self.load(path)
def Upsample(scale_factor=None, mode='nearest', align_corners=False):
return nn.Upsample(None, scale_factor, mode, align_corners)
def _make_fuse_layers(self, freeze_norm=False, norm_decay=0.):
if self.num_branches == 1:
return None
fuse_layers = []
num_out_branches = self.num_branches if self.multiscale_output else 1
for i in range(num_out_branches):
fuse_layer = []
for j in range(self.num_branches):
if j > i:
fuse_layer.append(
nn.Sequential(
L.Conv2d(
self.in_channels[j],
self.in_channels[i],
kernel_size=1,
stride=1,
padding=0,
bias=False,
), nn.BatchNorm2D(self.in_channels[i]),
nn.Upsample(scale_factor=2**(j - i),
mode='nearest')))
elif j == i:
fuse_layer.append(None)
else:
conv_downsamples = []
for k in range(i - j):
if k == i - j - 1:
conv_downsamples.append(
nn.Sequential(
L.Conv2d(
self.in_channels[j],
self.in_channels[j],
kernel_size=3,
stride=2,
padding=1,
groups=self.in_channels[j],
bias=False,
), nn.BatchNorm2D(self.in_channels[j]),
L.Conv2d(
self.in_channels[j],
self.in_channels[i],
kernel_size=1,
stride=1,
padding=0,
bias=False,
), nn.BatchNorm2D(self.in_channels[i])))
else:
conv_downsamples.append(
nn.Sequential(
L.Conv2d(
self.in_channels[j],
self.in_channels[j],
kernel_size=3,
stride=2,
padding=1,
groups=self.in_channels[j],
bias=False,
), nn.BatchNorm2D(self.in_channels[j]),
L.Conv2d(
self.in_channels[j],
self.in_channels[j],
kernel_size=1,
stride=1,
padding=0,
bias=False,
), nn.BatchNorm2D(self.in_channels[j]),
nn.ReLU()))
fuse_layer.append(nn.Sequential(*conv_downsamples))
fuse_layers.append(nn.LayerList(fuse_layer))
return nn.LayerList(fuse_layers)
def __init__(self,
in_channels=256,
mid_channels=128,
out_channels=256,
start_level=0,
end_level=3,
use_dcn_in_tower=False):
super(SOLOv2MaskHead, self).__init__()
assert start_level >= 0 and end_level >= start_level
self.in_channels = in_channels
self.out_channels = out_channels
self.mid_channels = mid_channels
self.use_dcn_in_tower = use_dcn_in_tower
self.range_level = end_level - start_level + 1
# TODO: add DeformConvNorm
conv_type = [ConvNormLayer]
self.conv_func = conv_type[0]
if self.use_dcn_in_tower:
self.conv_func = conv_type[1]
self.convs_all_levels = []
for i in range(start_level, end_level + 1):
conv_feat_name = 'mask_feat_head.convs_all_levels.{}'.format(i)
conv_pre_feat = nn.Sequential()
if i == start_level:
conv_pre_feat.add_sublayer(
conv_feat_name + '.conv' + str(i),
self.conv_func(ch_in=self.in_channels,
ch_out=self.mid_channels,
filter_size=3,
stride=1,
norm_type='gn',
norm_name=conv_feat_name + '.conv' +
str(i) + '.gn',
name=conv_feat_name + '.conv' + str(i)))
self.add_sublayer('conv_pre_feat' + str(i), conv_pre_feat)
self.convs_all_levels.append(conv_pre_feat)
else:
for j in range(i):
ch_in = 0
if j == 0:
ch_in = self.in_channels + 2 if i == end_level else self.in_channels
else:
ch_in = self.mid_channels
conv_pre_feat.add_sublayer(
conv_feat_name + '.conv' + str(j),
self.conv_func(ch_in=ch_in,
ch_out=self.mid_channels,
filter_size=3,
stride=1,
norm_type='gn',
norm_name=conv_feat_name + '.conv' +
str(j) + '.gn',
name=conv_feat_name + '.conv' + str(j)))
conv_pre_feat.add_sublayer(
conv_feat_name + '.conv' + str(j) + 'act', nn.ReLU())
conv_pre_feat.add_sublayer(
'upsample' + str(i) + str(j),
nn.Upsample(scale_factor=2, mode='bilinear'))
self.add_sublayer('conv_pre_feat' + str(i), conv_pre_feat)
self.convs_all_levels.append(conv_pre_feat)
conv_pred_name = 'mask_feat_head.conv_pred.0'
self.conv_pred = self.add_sublayer(
conv_pred_name,
self.conv_func(ch_in=self.mid_channels,
ch_out=self.out_channels,
filter_size=1,
stride=1,
norm_type='gn',
norm_name=conv_pred_name + '.gn',
name=conv_pred_name))
def __init__(self,
input_nc,
output_nc,
ngf=64,
n_blocks=6,
img_size=256,
light=False,
norm_type='instance'):
assert (n_blocks >= 0)
super(ResnetUGATITGenerator, self).__init__()
self.input_nc = input_nc
self.output_nc = output_nc
self.ngf = ngf
self.n_blocks = n_blocks
self.img_size = img_size
self.light = light
norm_layer = build_norm_layer(norm_type)
DownBlock = []
DownBlock += [
nn.Pad2D(padding=[3, 3, 3, 3], mode="reflect"),
nn.Conv2D(input_nc,
ngf,
kernel_size=7,
stride=1,
padding=0,
bias_attr=False),
norm_layer(ngf),
nn.ReLU()
]
# Down-Sampling
n_downsampling = 2
for i in range(n_downsampling):
mult = 2**i
DownBlock += [
nn.Pad2D(padding=[1, 1, 1, 1], mode="reflect"),
nn.Conv2D(ngf * mult,
ngf * mult * 2,
kernel_size=3,
stride=2,
padding=0,
bias_attr=False),
norm_layer(ngf * mult * 2),
nn.ReLU()
]
# Down-Sampling Bottleneck
mult = 2**n_downsampling
for i in range(n_blocks):
DownBlock += [
ResnetBlock(ngf * mult, use_bias=False, norm_layer=norm_layer)
]
# Class Activation Map
self.gap_fc = nn.Linear(ngf * mult, 1, bias_attr=False)
self.gmp_fc = nn.Linear(ngf * mult, 1, bias_attr=False)
self.conv1x1 = nn.Conv2D(ngf * mult * 2,
ngf * mult,
kernel_size=1,
stride=1,
bias_attr=True)
self.relu = nn.ReLU()
# Gamma, Beta block
if self.light:
FC = [
nn.Linear(ngf * mult, ngf * mult, bias_attr=False),
nn.ReLU(),
nn.Linear(ngf * mult, ngf * mult, bias_attr=False),
nn.ReLU()
]
else:
FC = [
nn.Linear(img_size // mult * img_size // mult * ngf * mult,
ngf * mult,
bias_attr=False),
nn.ReLU(),
nn.Linear(ngf * mult, ngf * mult, bias_attr=False),
nn.ReLU()
]
self.gamma = nn.Linear(ngf * mult, ngf * mult, bias_attr=False)
self.beta = nn.Linear(ngf * mult, ngf * mult, bias_attr=False)
# Up-Sampling Bottleneck
for i in range(n_blocks):
setattr(self, 'UpBlock1_' + str(i + 1),
ResnetAdaILNBlock(ngf * mult, use_bias=False))
# Up-Sampling
UpBlock2 = []
for i in range(n_downsampling):
mult = 2**(n_downsampling - i)
UpBlock2 += [
nn.Upsample(scale_factor=2, mode='nearest'),
nn.Pad2D(padding=[1, 1, 1, 1], mode="reflect"),
nn.Conv2D(ngf * mult,
int(ngf * mult / 2),
kernel_size=3,
stride=1,
padding=0,
bias_attr=False),
ILN(int(ngf * mult / 2)),
nn.ReLU()
]
UpBlock2 += [
nn.Pad2D(padding=[3, 3, 3, 3], mode="reflect"),
nn.Conv2D(ngf,
output_nc,
kernel_size=7,
stride=1,
padding=0,
bias_attr=False),
nn.Tanh()
]
self.DownBlock = nn.Sequential(*DownBlock)
self.FC = nn.Sequential(*FC)
self.UpBlock2 = nn.Sequential(*UpBlock2)
def _make_fuse_layers(self):
if self.num_branches == 1:
return None
num_branches = self.num_branches
num_inchannels = self.num_inchannels
fuse_layers = []
for i in range(num_branches if self.multi_scale_output else 1):
fuse_layer = []
for j in range(num_branches):
if j > i:
fuse_layer.append(
nn.Sequential(
# nn.Conv2d(
# num_inchannels[j],
# num_inchannels[i],
# 1, 1, 0, bias=False
# ),
nn.Conv2D(num_inchannels[j],
num_inchannels[i],
1,
1,
0,
weight_attr=weight_attr_conv,
bias_attr=False),
# nn.BatchNorm2d(num_inchannels[i]),
nn.BatchNorm2D(num_inchannels[i],
weight_attr=weight_attr_bn,
bias_attr=bias_attr_constant_0),
nn.Upsample(scale_factor=2**(j - i),
mode='nearest')))
elif j == i:
# fuse_layer.append(None)
fuse_layer.append(equalmap())
else:
conv3x3s = []
for k in range(i - j):
if k == i - j - 1:
num_outchannels_conv3x3 = num_inchannels[i]
conv3x3s.append(
nn.Sequential(
# nn.Conv2d(
# num_inchannels[j],
# num_outchannels_conv3x3,
# 3, 2, 1, bias=False
# ),
nn.Conv2D(num_inchannels[j],
num_outchannels_conv3x3,
3,
2,
1,
weight_attr=weight_attr_conv,
bias_attr=False),
# nn.BatchNorm2d(num_outchannels_conv3x3)
nn.BatchNorm2D(
num_outchannels_conv3x3,
weight_attr=weight_attr_bn,
bias_attr=bias_attr_constant_0)))
else:
num_outchannels_conv3x3 = num_inchannels[j]
conv3x3s.append(
nn.Sequential(
# nn.Conv2d(
# num_inchannels[j],
# num_outchannels_conv3x3,
# 3, 2, 1, bias=False
# ),
nn.Conv2D(num_inchannels[j],
num_outchannels_conv3x3,
3,
2,
1,
weight_attr=weight_attr_conv,
bias_attr=False),
# nn.BatchNorm2d(num_outchannels_conv3x3),
nn.BatchNorm2D(
num_outchannels_conv3x3,
weight_attr=weight_attr_bn,
bias_attr=bias_attr_constant_0),
# nn.ReLU(True)
nn.ReLU()))
fuse_layer.append(nn.Sequential(*conv3x3s))
# fuse_layers.append(nn.ModuleList(fuse_layer))
fuse_layers.append(nn.LayerList(fuse_layer))
# return nn.ModuleList(fuse_layers)
return nn.LayerList(fuse_layers)
def __init__(self,
backbone,
pretrained=None,
backbone_scale=0.25,
refine_kernel_size=3,
if_refine=True):
super().__init__()
if if_refine:
if backbone_scale > 0.5:
raise ValueError(
'Backbone_scale should not be greater than 1/2, but it is {}'
.format(backbone_scale))
else:
backbone_scale = 1
self.backbone = backbone
self.backbone_scale = backbone_scale
self.pretrained = pretrained
self.if_refine = if_refine
if if_refine:
self.refiner = Refiner(kernel_size=refine_kernel_size)
self.backbone_channels = backbone.feat_channels
######################
### Decoder part - Glance
######################
self.psp_module = layers.PPModule(self.backbone_channels[-1],
512,
bin_sizes=(1, 3, 5),
dim_reduction=False,
align_corners=False)
self.psp4 = conv_up_psp(512, 256, 2)
self.psp3 = conv_up_psp(512, 128, 4)
self.psp2 = conv_up_psp(512, 64, 8)
self.psp1 = conv_up_psp(512, 64, 16)
# stage 5g
self.decoder5_g = nn.Sequential(
layers.ConvBNReLU(512 + self.backbone_channels[-1],
512,
3,
padding=1),
layers.ConvBNReLU(512, 512, 3, padding=2, dilation=2),
layers.ConvBNReLU(512, 256, 3, padding=2, dilation=2),
nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False))
# stage 4g
self.decoder4_g = nn.Sequential(
layers.ConvBNReLU(512, 256, 3, padding=1),
layers.ConvBNReLU(256, 256, 3, padding=1),
layers.ConvBNReLU(256, 128, 3, padding=1),
nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False))
# stage 3g
self.decoder3_g = nn.Sequential(
layers.ConvBNReLU(256, 128, 3, padding=1),
layers.ConvBNReLU(128, 128, 3, padding=1),
layers.ConvBNReLU(128, 64, 3, padding=1),
nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False))
# stage 2g
self.decoder2_g = nn.Sequential(
layers.ConvBNReLU(128, 128, 3, padding=1),
layers.ConvBNReLU(128, 128, 3, padding=1),
layers.ConvBNReLU(128, 64, 3, padding=1),
nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False))
# stage 1g
self.decoder1_g = nn.Sequential(
layers.ConvBNReLU(128, 64, 3, padding=1),
layers.ConvBNReLU(64, 64, 3, padding=1),
layers.ConvBNReLU(64, 64, 3, padding=1),
nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False))
# stage 0g
self.decoder0_g = nn.Sequential(
layers.ConvBNReLU(64, 64, 3, padding=1),
layers.ConvBNReLU(64, 64, 3, padding=1),
nn.Conv2D(64, 3, 3, padding=1))
##########################
### Decoder part - FOCUS
##########################
self.bridge_block = nn.Sequential(
layers.ConvBNReLU(self.backbone_channels[-1],
512,
3,
dilation=2,
padding=2),
layers.ConvBNReLU(512, 512, 3, dilation=2, padding=2),
layers.ConvBNReLU(512, 512, 3, dilation=2, padding=2))
# stage 5f
self.decoder5_f = nn.Sequential(
layers.ConvBNReLU(512 + self.backbone_channels[-1],
512,
3,
padding=1),
layers.ConvBNReLU(512, 512, 3, padding=2, dilation=2),
layers.ConvBNReLU(512, 256, 3, padding=2, dilation=2),
nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False))
# stage 4f
self.decoder4_f = nn.Sequential(
layers.ConvBNReLU(256 + self.backbone_channels[-2],
256,
3,
padding=1),
layers.ConvBNReLU(256, 256, 3, padding=1),
layers.ConvBNReLU(256, 128, 3, padding=1),
nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False))
# stage 3f
self.decoder3_f = nn.Sequential(
layers.ConvBNReLU(128 + self.backbone_channels[-3],
128,
3,
padding=1),
layers.ConvBNReLU(128, 128, 3, padding=1),
layers.ConvBNReLU(128, 64, 3, padding=1),
nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False))
# stage 2f
self.decoder2_f = nn.Sequential(
layers.ConvBNReLU(64 + self.backbone_channels[-4],
128,
3,
padding=1),
layers.ConvBNReLU(128, 128, 3, padding=1),
layers.ConvBNReLU(128, 64, 3, padding=1),
nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False))
# stage 1f
self.decoder1_f = nn.Sequential(
layers.ConvBNReLU(64 + self.backbone_channels[-5],
64,
3,
padding=1),
layers.ConvBNReLU(64, 64, 3, padding=1),
layers.ConvBNReLU(64, 64, 3, padding=1),
nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False))
# stage 0f
self.decoder0_f = nn.Sequential(
layers.ConvBNReLU(64, 64, 3, padding=1),
layers.ConvBNReLU(64, 64, 3, padding=1),
nn.Conv2D(64, 1 + 1 + 32, 3, padding=1))
self.init_weight()
def __init__(self,
in_nf=3,
out_nf=3,
scale_factor=4,
nf=64,
nframes=5,
groups=8,
front_RBs=5,
back_RBs=10,
center=None,
predeblur=False,
HR_in=False,
w_TSA=True,
TSA_only=False):
super(EDVRNet, self).__init__()
self.in_nf = in_nf
self.out_nf = out_nf
self.scale_factor = scale_factor
self.nf = nf
self.nframes = nframes
self.groups = groups
self.front_RBs = front_RBs
self.back_RBs = back_RBs
self.center = nframes // 2 if center is None else center
self.predeblur = True if predeblur else False
self.HR_in = True if HR_in else False
self.w_TSA = True if w_TSA else False
self.Leaky_relu = nn.LeakyReLU(negative_slope=0.1)
if self.predeblur:
self.pre_deblur = PredeblurResNetPyramid(in_nf=self.in_nf,
nf=self.nf,
HR_in=self.HR_in)
self.cov_1 = nn.Conv2D(in_channels=self.nf,
out_channels=self.nf,
kernel_size=1,
stride=1)
else:
if self.HR_in:
self.conv_first_1 = nn.Conv2D(in_channels=self.in_nf,
out_channels=self.nf,
kernel_size=3,
stride=1,
padding=1)
self.conv_first_2 = nn.Conv2D(in_channels=self.nf,
out_channels=self.nf,
kernel_size=3,
stride=2,
padding=1)
self.conv_first_3 = nn.Conv2D(in_channels=self.nf,
out_channels=self.nf,
kernel_size=3,
stride=2,
padding=1)
else:
self.conv_first = nn.Conv2D(in_channels=self.in_nf,
out_channels=self.nf,
kernel_size=3,
stride=1,
padding=1)
#feature extraction module
self.feature_extractor = MakeMultiBlocks(ResidualBlockNoBN,
self.front_RBs, self.nf)
self.fea_L2_conv1 = nn.Conv2D(in_channels=self.nf,
out_channels=self.nf,
kernel_size=3,
stride=2,
padding=1)
self.fea_L2_conv2 = nn.Conv2D(in_channels=self.nf,
out_channels=self.nf,
kernel_size=3,
stride=1,
padding=1)
self.fea_L3_conv1 = nn.Conv2D(
in_channels=self.nf,
out_channels=self.nf,
kernel_size=3,
stride=2,
padding=1,
)
self.fea_L3_conv2 = nn.Conv2D(in_channels=self.nf,
out_channels=self.nf,
kernel_size=3,
stride=1,
padding=1)
#PCD alignment module
self.PCDModule = PCDAlign(nf=self.nf, groups=self.groups)
#TSA Fusion module
if self.w_TSA:
self.TSAModule = TSAFusion(nf=self.nf,
nframes=self.nframes,
center=self.center)
else:
self.TSAModule = nn.Conv2D(in_channels=self.nframes * self.nf,
out_channels=self.nf,
kernel_size=1,
stride=1)
#reconstruction module
self.reconstructor = MakeMultiBlocks(ResidualBlockNoBN, self.back_RBs,
self.nf)
self.upconv1 = nn.Conv2D(in_channels=self.nf,
out_channels=4 * self.nf,
kernel_size=3,
stride=1,
padding=1)
self.pixel_shuffle = nn.PixelShuffle(2)
self.upconv2 = nn.Conv2D(in_channels=self.nf,
out_channels=4 * self.nf,
kernel_size=3,
stride=1,
padding=1)
self.HRconv = nn.Conv2D(in_channels=self.nf,
out_channels=self.nf,
kernel_size=3,
stride=1,
padding=1)
self.conv_last = nn.Conv2D(in_channels=self.nf,
out_channels=self.out_nf,
kernel_size=3,
stride=1,
padding=1)
self.upsample = nn.Upsample(scale_factor=self.scale_factor,
mode="bilinear",
align_corners=False,
align_mode=0)
def __init__(self, nf=64, nframes=5, center=2):
super(TSAFusion, self).__init__()
self.nf = nf
self.nframes = nframes
self.center = center
self.sigmoid = nn.Sigmoid()
self.Leaky_relu = nn.LeakyReLU(negative_slope=0.1)
self.tAtt_2 = nn.Conv2D(in_channels=self.nf,
out_channels=self.nf,
kernel_size=3,
stride=1,
padding=1)
self.tAtt_1 = nn.Conv2D(in_channels=self.nf,
out_channels=self.nf,
kernel_size=3,
stride=1,
padding=1)
self.fea_fusion = nn.Conv2D(in_channels=self.nf * self.nframes,
out_channels=self.nf,
kernel_size=1,
stride=1,
padding=0)
self.sAtt_1 = nn.Conv2D(in_channels=self.nf * self.nframes,
out_channels=self.nf,
kernel_size=1,
stride=1,
padding=0)
self.max_pool = nn.MaxPool2D(3, stride=2, padding=1)
self.avg_pool = nn.AvgPool2D(3, stride=2, padding=1, exclusive=False)
self.sAtt_2 = nn.Conv2D(in_channels=2 * self.nf,
out_channels=self.nf,
kernel_size=1,
stride=1,
padding=0)
self.sAtt_3 = nn.Conv2D(in_channels=self.nf,
out_channels=self.nf,
kernel_size=3,
stride=1,
padding=1)
self.sAtt_4 = nn.Conv2D(
in_channels=self.nf,
out_channels=self.nf,
kernel_size=1,
stride=1,
padding=0,
)
self.sAtt_5 = nn.Conv2D(in_channels=self.nf,
out_channels=self.nf,
kernel_size=3,
stride=1,
padding=1)
self.sAtt_add_1 = nn.Conv2D(in_channels=self.nf,
out_channels=self.nf,
kernel_size=1,
stride=1,
padding=0)
self.sAtt_add_2 = nn.Conv2D(in_channels=self.nf,
out_channels=self.nf,
kernel_size=1,
stride=1,
padding=0)
self.sAtt_L1 = nn.Conv2D(in_channels=self.nf,
out_channels=self.nf,
kernel_size=1,
stride=1,
padding=0)
self.sAtt_L2 = nn.Conv2D(
in_channels=2 * self.nf,
out_channels=self.nf,
kernel_size=3,
stride=1,
padding=1,
)
self.sAtt_L3 = nn.Conv2D(in_channels=self.nf,
out_channels=self.nf,
kernel_size=3,
stride=1,
padding=1)
self.upsample = nn.Upsample(scale_factor=2,
mode="bilinear",
align_corners=False,
align_mode=0)
def __init__(self, in_channels, out_channels, scale_factor):
super(ConvUp2D, self).__init__(
nn.Conv2D(in_channels, out_channels, 3, padding=1),
nn.Upsample(scale_factor=scale_factor, mode='bilinear'))
def __init__(self, filters, cat_channels, up_channels):
super(Decoder, self).__init__()
'''stage 4d'''
# h1->320*320, hd4->40*40, Pooling 8 times
self.h1_PT_hd4 = nn.MaxPool2D(8, 8, ceil_mode=True)
self.h1_PT_hd4_cbr = ConvBnReLU2D(filters[0], cat_channels)
# h2->160*160, hd4->40*40, Pooling 4 times
self.h2_PT_hd4 = nn.MaxPool2D(4, 4, ceil_mode=True)
self.h2_PT_hd4_cbr = ConvBnReLU2D(filters[1], cat_channels)
# h3->80*80, hd4->40*40, Pooling 2 times
self.h3_PT_hd4 = nn.MaxPool2D(2, 2, ceil_mode=True)
self.h3_PT_hd4_cbr = ConvBnReLU2D(filters[2], cat_channels)
# h4->40*40, hd4->40*40, Concatenation
self.h4_Cat_hd4_cbr = ConvBnReLU2D(filters[3], cat_channels)
# hd5->20*20, hd4->40*40, Upsample 2 times
self.hd5_UT_hd4 = nn.Upsample(scale_factor=2, mode='bilinear') # 14*14
self.hd5_UT_hd4_cbr = ConvBnReLU2D(filters[4], cat_channels)
# fusion(h1_PT_hd4, h2_PT_hd4, h3_PT_hd4, h4_Cat_hd4, hd5_UT_hd4)
self.cbr4d_1 = ConvBnReLU2D(up_channels, up_channels) # 16
'''stage 3d'''
# h1->320*320, hd3->80*80, Pooling 4 times
self.h1_PT_hd3 = nn.MaxPool2D(4, 4, ceil_mode=True)
self.h1_PT_hd3_cbr = ConvBnReLU2D(filters[0], cat_channels)
# h2->160*160, hd3->80*80, Pooling 2 times
self.h2_PT_hd3 = nn.MaxPool2D(2, 2, ceil_mode=True)
self.h2_PT_hd3_cbr = ConvBnReLU2D(filters[1], cat_channels)
# h3->80*80, hd3->80*80, Concatenation
self.h3_Cat_hd3_cbr = ConvBnReLU2D(filters[2], cat_channels)
# hd4->40*40, hd4->80*80, Upsample 2 times
self.hd4_UT_hd3 = nn.Upsample(scale_factor=2, mode='bilinear') # 14*14
self.hd4_UT_hd3_cbr = ConvBnReLU2D(up_channels, cat_channels)
# hd5->20*20, hd4->80*80, Upsample 4 times
self.hd5_UT_hd3 = nn.Upsample(scale_factor=4, mode='bilinear') # 14*14
self.hd5_UT_hd3_cbr = ConvBnReLU2D(filters[4], cat_channels)
# fusion(h1_PT_hd3, h2_PT_hd3, h3_Cat_hd3, hd4_UT_hd3, hd5_UT_hd3)
self.cbr3d_1 = ConvBnReLU2D(up_channels, up_channels) # 16
'''stage 2d '''
# h1->320*320, hd2->160*160, Pooling 2 times
self.h1_PT_hd2 = nn.MaxPool2D(2, 2, ceil_mode=True)
self.h1_PT_hd2_cbr = ConvBnReLU2D(filters[0], cat_channels)
# h2->160*160, hd2->160*160, Concatenation
self.h2_Cat_hd2_cbr = ConvBnReLU2D(filters[1], cat_channels)
# hd3->80*80, hd2->160*160, Upsample 2 times
self.hd3_UT_hd2 = nn.Upsample(scale_factor=2, mode='bilinear') # 14*14
self.hd3_UT_hd2_cbr = ConvBnReLU2D(up_channels, cat_channels)
# hd4->40*40, hd2->160*160, Upsample 4 times
self.hd4_UT_hd2 = nn.Upsample(scale_factor=4, mode='bilinear') # 14*14
self.hd4_UT_hd2_cbr = ConvBnReLU2D(up_channels, cat_channels)
# hd5->20*20, hd2->160*160, Upsample 8 times
self.hd5_UT_hd2 = nn.Upsample(scale_factor=8, mode='bilinear') # 14*14
self.hd5_UT_hd2_cbr = ConvBnReLU2D(filters[4], cat_channels)
# fusion(h1_PT_hd2, h2_Cat_hd2, hd3_UT_hd2, hd4_UT_hd2, hd5_UT_hd2)
self.cbr2d_1 = ConvBnReLU2D(up_channels, up_channels) # 16
'''stage 1d'''
# h1->320*320, hd1->320*320, Concatenation
self.h1_Cat_hd1_cbr = ConvBnReLU2D(filters[0], cat_channels)
# hd2->160*160, hd1->320*320, Upsample 2 times
self.hd2_UT_hd1 = nn.Upsample(scale_factor=2, mode='bilinear') # 14*14
self.hd2_UT_hd1_cbr = ConvBnReLU2D(up_channels, cat_channels)
# hd3->80*80, hd1->320*320, Upsample 4 times
self.hd3_UT_hd1 = nn.Upsample(scale_factor=4, mode='bilinear') # 14*14
self.hd3_UT_hd1_cbr = ConvBnReLU2D(up_channels, cat_channels)
# hd4->40*40, hd1->320*320, Upsample 8 times
self.hd4_UT_hd1 = nn.Upsample(scale_factor=8, mode='bilinear') # 14*14
self.hd4_UT_hd1_cbr = ConvBnReLU2D(up_channels, cat_channels)
# hd5->20*20, hd1->320*320, Upsample 16 times
self.hd5_UT_hd1 = nn.Upsample(scale_factor=16, mode='bilinear') # 14*14
self.hd5_UT_hd1_cbr = ConvBnReLU2D(filters[4], cat_channels)
# fusion(h1_Cat_hd1, hd2_UT_hd1, hd3_UT_hd1, hd4_UT_hd1, hd5_UT_hd1)
self.cbr1d_1 = ConvBnReLU2D(up_channels, up_channels) # 16
def __init__(self, ch_in, ch_out):
super().__init__()
self.up = nn.Sequential(
nn.Upsample(scale_factor=2, mode="bilinear"),
nn.Conv2D(ch_in, ch_out, kernel_size=3, stride=1, padding=1),
nn.BatchNorm2D(ch_out), nn.ReLU())
def __init__(self, nf=64, groups=8):
super(PCDAlign, self).__init__()
self.nf = nf
self.groups = groups
self.Leaky_relu = nn.LeakyReLU(negative_slope=0.1)
self.upsample = nn.Upsample(scale_factor=2,
mode="bilinear",
align_corners=False,
align_mode=0)
# Pyramid has three levels:
# L3: level 3, 1/4 spatial size
# L2: level 2, 1/2 spatial size
# L1: level 1, original spatial size
# L3
self.PCD_Align_L3_offset_conv1 = nn.Conv2D(in_channels=nf * 2,
out_channels=nf,
kernel_size=3,
stride=1,
padding=1)
self.PCD_Align_L3_offset_conv2 = nn.Conv2D(in_channels=nf,
out_channels=nf,
kernel_size=3,
stride=1,
padding=1)
self.PCD_Align_L3_dcn = DCNPack(num_filters=nf,
kernel_size=3,
stride=1,
padding=1,
deformable_groups=groups)
#L2
self.PCD_Align_L2_offset_conv1 = nn.Conv2D(in_channels=nf * 2,
out_channels=nf,
kernel_size=3,
stride=1,
padding=1)
self.PCD_Align_L2_offset_conv2 = nn.Conv2D(in_channels=nf * 2,
out_channels=nf,
kernel_size=3,
stride=1,
padding=1)
self.PCD_Align_L2_offset_conv3 = nn.Conv2D(in_channels=nf,
out_channels=nf,
kernel_size=3,
stride=1,
padding=1)
self.PCD_Align_L2_dcn = DCNPack(num_filters=nf,
kernel_size=3,
stride=1,
padding=1,
deformable_groups=groups)
self.PCD_Align_L2_fea_conv = nn.Conv2D(in_channels=nf * 2,
out_channels=nf,
kernel_size=3,
stride=1,
padding=1)
#L1
self.PCD_Align_L1_offset_conv1 = nn.Conv2D(in_channels=nf * 2,
out_channels=nf,
kernel_size=3,
stride=1,
padding=1)
self.PCD_Align_L1_offset_conv2 = nn.Conv2D(in_channels=nf * 2,
out_channels=nf,
kernel_size=3,
stride=1,
padding=1)
self.PCD_Align_L1_offset_conv3 = nn.Conv2D(in_channels=nf,
out_channels=nf,
kernel_size=3,
stride=1,
padding=1)
self.PCD_Align_L1_dcn = DCNPack(num_filters=nf,
kernel_size=3,
stride=1,
padding=1,
deformable_groups=groups)
self.PCD_Align_L1_fea_conv = nn.Conv2D(in_channels=nf * 2,
out_channels=nf,
kernel_size=3,
stride=1,
padding=1)
#cascade
self.PCD_Align_cas_offset_conv1 = nn.Conv2D(in_channels=nf * 2,
out_channels=nf,
kernel_size=3,
stride=1,
padding=1)
self.PCD_Align_cas_offset_conv2 = nn.Conv2D(in_channels=nf,
out_channels=nf,
kernel_size=3,
stride=1,
padding=1)
self.PCD_Align_cascade_dcn = DCNPack(num_filters=nf,
kernel_size=3,
stride=1,
padding=1,
deformable_groups=groups)
def __init__(self,
num_class,
fpn_inplanes,
channels,
dropout_ratio=0.1,
fpn_dim=256,
enable_auxiliary_loss=False,
align_corners=False):
super(PFPNHead, self).__init__()
self.enable_auxiliary_loss = enable_auxiliary_loss
self.align_corners = align_corners
self.lateral_convs = nn.LayerList()
self.fpn_out = nn.LayerList()
for fpn_inplane in fpn_inplanes:
self.lateral_convs.append(
nn.Sequential(nn.Conv2D(fpn_inplane, fpn_dim, 1),
layers.SyncBatchNorm(fpn_dim), nn.ReLU()))
self.fpn_out.append(
nn.Sequential(
layers.ConvBNReLU(fpn_dim, fpn_dim, 3, bias_attr=False)))
self.scale_heads = nn.LayerList()
for index in range(len(fpn_inplanes)):
head_length = max(
1,
int(np.log2(fpn_inplanes[index]) - np.log2(fpn_inplanes[0])))
scale_head = nn.LayerList()
for head_index in range(head_length):
scale_head.append(
layers.ConvBNReLU(
fpn_dim,
channels,
3,
padding=1,
))
if fpn_inplanes[index] != fpn_inplanes[0]:
scale_head.append(
nn.Upsample(scale_factor=2,
mode='bilinear',
align_corners=align_corners))
self.scale_heads.append(nn.Sequential(*scale_head))
if dropout_ratio:
self.dropout = nn.Dropout2D(dropout_ratio)
if self.enable_auxiliary_loss:
self.dsn = nn.Sequential(
layers.ConvBNReLU(fpn_inplanes[2],
fpn_inplanes[2],
3,
padding=1), nn.Dropout2D(dropout_ratio),
nn.Conv2D(fpn_inplanes[2], num_class, kernel_size=1))
else:
self.dropout = None
if self.enable_auxiliary_loss:
self.dsn = nn.Sequential(
layers.ConvBNReLU(fpn_inplanes[2],
fpn_inplanes[2],
3,
padding=1),
nn.Conv2D(fpn_inplanes[2], num_class, kernel_size=1))
self.conv_last = nn.Sequential(
layers.ConvBNReLU(len(fpn_inplanes) * fpn_dim,
fpn_dim,
3,
bias_attr=False),
nn.Conv2D(fpn_dim, num_class, kernel_size=1))
self.conv_seg = nn.Conv2D(channels, num_class, kernel_size=1)
def conv_up_psp(in_channels, out_channels, up_sample):
return nn.Sequential(
layers.ConvBNReLU(in_channels, out_channels, 3, padding=1),
nn.Upsample(scale_factor=up_sample,
mode='bilinear',
align_corners=False))
def my_upsample(in_channels,out_channels,in_size,multiple,type='transposeconv',upsample_type='nearest'):
if type == 'transposeconv':
return nn.Conv2DTranspose(in_channels,out_channels,4,2,1)
if type == 'upsample':
return nn.Upsample(size=[in_size*multiple,in_size*multiple],mode=upsample_type)
def __init__(self, ngf=32, img_size=256, n_blocks=4, light=True):
super(ResnetGenerator, self).__init__()
self.light = light
self.n_blocks = n_blocks
DownBlock = []
DownBlock += [
nn.Pad2D([3, 3, 3, 3], 'reflect'),
nn.Conv2D(3, ngf, kernel_size=7, stride=1, bias_attr=False),
nn.InstanceNorm2D(ngf, weight_attr=False, bias_attr=False),
nn.ReLU()
]
DownBlock += [HourGlass(ngf, ngf), HourGlass(ngf, ngf)]
# Down-Sampling
n_downsampling = 2
for i in range(n_downsampling):
mult = 2**i
DownBlock += [
nn.Pad2D([1, 1, 1, 1], 'reflect'),
nn.Conv2D(ngf * mult,
ngf * mult * 2,
kernel_size=3,
stride=2,
bias_attr=False),
nn.InstanceNorm2D(ngf * mult * 2,
weight_attr=False,
bias_attr=False),
nn.ReLU()
]
# Encoder Bottleneck
mult = 2**n_downsampling
for i in range(n_blocks):
setattr(self, 'EncodeBlock' + str(i + 1), ResnetBlock(ngf * mult))
# Class Activation Map
self.gap_fc = nn.Linear(ngf * mult, 1, bias_attr=False)
self.gmp_fc = nn.Linear(ngf * mult, 1, bias_attr=False)
self.conv1x1 = nn.Conv2D(ngf * mult * 2,
ngf * mult,
kernel_size=1,
stride=1)
self.relu = nn.ReLU()
# Gamma, Beta block
FC = []
if self.light:
FC += [
nn.Linear(ngf * mult, ngf * mult, bias_attr=False),
nn.ReLU(),
nn.Linear(ngf * mult, ngf * mult, bias_attr=False),
nn.ReLU()
]
else:
FC += [
nn.Linear(img_size // mult * img_size // mult * ngf * mult,
ngf * mult,
bias_attr=False),
nn.ReLU(),
nn.Linear(ngf * mult, ngf * mult, bias_attr=False),
nn.ReLU()
]
# Decoder Bottleneck
mult = 2**n_downsampling
for i in range(n_blocks):
setattr(self, 'DecodeBlock' + str(i + 1),
ResnetSoftAdaLINBlock(ngf * mult))
# Up-Sampling
UpBlock = []
for i in range(n_downsampling):
mult = 2**(n_downsampling - i)
UpBlock += [
nn.Upsample(scale_factor=2),
nn.Pad2D([1, 1, 1, 1], 'reflect'),
nn.Conv2D(ngf * mult,
ngf * mult // 2,
kernel_size=3,
stride=1,
bias_attr=False),
LIN(ngf * mult // 2),
nn.ReLU()
]
UpBlock += [HourGlass(ngf, ngf), HourGlass(ngf, ngf, False)]
UpBlock += [
nn.Pad2D([3, 3, 3, 3], 'reflect'),
nn.Conv2D(3, 3, kernel_size=7, stride=1, bias_attr=False),
nn.Tanh()
]
self.DownBlock = nn.Sequential(*DownBlock)
self.FC = nn.Sequential(*FC)
self.UpBlock = nn.Sequential(*UpBlock)
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)'
)
