def __init__(self, out, eps, momentum):
super(LayerTest, self).__init__()
self.in2d = nn.InstanceNorm2d(out, eps=eps, momentum=momentum)
def get_normalization(name, channels):
if name == 'in': return nn.InstanceNorm2d(channels)
if name == 'bn': return nn.BatchNorm2d(channels)
def __init__(self, hyperparameters):
super(aclmaskpermgidtno_Trainer, self).__init__()
lr = hyperparameters['lr']
# Initiate the networks
self.gen_AB = AdaINGen(
hyperparameters['input_dim_a'],
hyperparameters['gen']) # auto-encoder for domain A
self.gen_BA = AdaINGen(
hyperparameters['input_dim_a'],
hyperparameters['gen']) # auto-encoder for domain B
self.dis_A = MsImageDis(
hyperparameters['input_dim_a'],
hyperparameters['dis']) # discriminator for domain A
self.dis_B = MsImageDis(
hyperparameters['input_dim_a'],
hyperparameters['dis']) # discriminator for domain B
self.dis_2 = MsImageDis(hyperparameters['input_dim_b'],
hyperparameters['dis']) # discriminator 2
# self.dis_2B = MsImageDis(hyperparameters['input_dim_a'], hyperparameters['dis']) # discriminator 2 for domain B
self.instancenorm = nn.InstanceNorm2d(512, affine=False)
self.style_dim = hyperparameters['gen']['style_dim']
# fix the noise used in sampling
display_size = int(hyperparameters['display_size'])
self.z_1 = torch.randn(display_size, self.style_dim, 1, 1).cuda()
self.z_2 = torch.randn(display_size, self.style_dim, 1, 1).cuda()
self.z_3 = torch.randn(display_size, self.style_dim, 1, 1).cuda()
# Setup the optimizers
beta1 = hyperparameters['beta1']
beta2 = hyperparameters['beta2']
dis_params = list(self.dis_A.parameters()) + list(
self.dis_B.parameters()) + list(self.dis_2.parameters())
gen_params = list(self.gen_AB.parameters()) + list(
self.gen_BA.parameters())
self.dis_opt = torch.optim.Adam(
[p for p in dis_params if p.requires_grad],
lr=lr,
betas=(beta1, beta2),
weight_decay=hyperparameters['weight_decay'])
self.gen_opt = torch.optim.Adam(
[p for p in gen_params if p.requires_grad],
lr=lr,
betas=(beta1, beta2),
weight_decay=hyperparameters['weight_decay'])
self.dis_scheduler = get_scheduler(self.dis_opt, hyperparameters)
self.gen_scheduler = get_scheduler(self.gen_opt, hyperparameters)
self.alpha = hyperparameters['alpha']
self.focus_lam = hyperparameters['focus_loss']
# Network weight initialization
self.apply(weights_init(hyperparameters['init']))
self.dis_A.apply(weights_init('gaussian'))
self.dis_B.apply(weights_init('gaussian'))
self.dis_2.apply(weights_init('gaussian'))
# Load VGG model if needed
if 'vgg_w' in hyperparameters.keys() and hyperparameters['vgg_w'] > 0:
self.vgg = load_vgg16(hyperparameters['vgg_model_path'] +
'/models')
self.vgg.eval()
for param in self.vgg.parameters():
param.requires_grad = False
def __init__(self,
in_channels,
out_channels,
kernel_size,
stride=1,
padding=0,
dilation=1,
pad_type='zero',
activation='lrelu',
norm='none',
sn=False):
super(Conv2dLayer, self).__init__()
# Initialize the padding scheme
if pad_type == 'reflect':
self.pad = nn.ReflectionPad2d(padding)
elif pad_type == 'replicate':
self.pad = nn.ReplicationPad2d(padding)
elif pad_type == 'zero':
self.pad = nn.ZeroPad2d(padding)
else:
assert 0, "Unsupported padding type: {}".format(pad_type)
# Initialize the normalization type
if norm == 'bn':
self.norm = nn.BatchNorm2d(out_channels)
elif norm == 'in':
self.norm = nn.InstanceNorm2d(out_channels)
elif norm == 'ln':
self.norm = LayerNorm(out_channels)
elif norm == 'none':
self.norm = None
else:
assert 0, "Unsupported normalization: {}".format(norm)
# Initialize the activation funtion
if activation == 'relu':
self.activation = nn.ReLU(inplace=True)
elif activation == 'lrelu':
self.activation = nn.LeakyReLU(0.2, inplace=True)
elif activation == 'prelu':
self.activation = nn.PReLU()
elif activation == 'selu':
self.activation = nn.SELU(inplace=True)
elif activation == 'tanh':
self.activation = nn.Tanh()
elif activation == 'sigmoid':
self.activation = nn.Sigmoid()
elif activation == 'none':
self.activation = None
else:
assert 0, "Unsupported activation: {}".format(activation)
# Initialize the convolution layers
if sn:
self.conv2d = SpectralNorm(
nn.Conv2d(in_channels,
out_channels,
kernel_size,
stride,
padding=0,
dilation=dilation))
else:
self.conv2d = nn.Conv2d(in_channels,
out_channels,
kernel_size,
stride,
padding=0,
dilation=dilation)
def __init__(self, num_classes=31):
super(AlexNet, self).__init__()
self.relu = nn.ReLU(inplace=True)
self.maxpool = nn.MaxPool2d(kernel_size=2, stride=2)
self.selayer3 = SELayer(384)
self.selayer4 = SELayer(256)
self.selayer5 = SELayer(256)
self.ad_bn_3 = nn.AD_BatchNorm2d(384)
self.ad_bn_4 = nn.AD_BatchNorm2d(256)
self.ad_bn_5 = nn.AD_BatchNorm2d(256)
self.block1 = nn.Sequential(
nn.Conv2d(3, 64, kernel_size=11, stride=4, padding=5),
#SELayer(64)
nn.InstanceNorm2d(64),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=2, stride=2),
)
self.block2 = nn.Sequential(
nn.Conv2d(64, 192, kernel_size=5, padding=2),
nn.InstanceNorm2d(192),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=2, stride=2),
)
self.block3 = nn.Sequential(
nn.Conv2d(192, 384, kernel_size=3, padding=1), )
# nn.BatchNorm2d(384),
# nn.ReLU(inplace=True),
#)
self.block4 = nn.Sequential(
nn.Conv2d(384, 256, kernel_size=3, padding=1), )
# nn.BatchNorm2d(256),
# nn.ReLU(inplace=True),
#)
self.block5 = nn.Sequential(
nn.Conv2d(256, 256, kernel_size=3, padding=1), )
# nn.BatchNorm2d(256),
# nn.ReLU(inplace=True),
# nn.MaxPool2d(kernel_size=2, stride=2),
#nn.BatchNorm2d(256),
#)
self.features = nn.Sequential(
nn.Conv2d(3, 64, kernel_size=11, stride=4, padding=5),
nn.InstanceNorm2d(64),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=2, stride=2),
nn.Conv2d(64, 192, kernel_size=5, padding=2),
nn.InstanceNorm2d(192),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=2, stride=2),
nn.Conv2d(192, 384, kernel_size=3, padding=1),
nn.BatchNorm2d(384),
nn.ReLU(inplace=True),
nn.Conv2d(384, 256, kernel_size=3, padding=1),
nn.BatchNorm2d(256),
nn.ReLU(inplace=True),
nn.Conv2d(256, 256, kernel_size=3, padding=1),
nn.BatchNorm2d(256),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=2, stride=2),
)
#self.copy_feature = copy.deepcopy(self.features)
#print(self.features.1)
self.bn = nn.BatchNorm1d(256)
self.ad_pool = nn.AdaptiveAvgPool2d(1)
self.classifier = nn.Linear(256, num_classes)
def __init__(self, z_dim=8):
super(Generator, self).__init__()
self.in_channels = 3
self.ngf = 64
self.out_channels = 3
self.down_1 = nn.Sequential(
nn.Conv2d(self.in_channels + z_dim,
self.ngf,
kernel_size=4,
stride=2,
padding=1,
bias=False), nn.LeakyReLU(0.2))
self.down_2 = nn.Sequential(
nn.Conv2d(self.ngf,
self.ngf * 2,
kernel_size=4,
stride=2,
padding=1,
bias=False), nn.InstanceNorm2d(self.ngf * 2,
affine=True),
nn.LeakyReLU(0.2))
self.down_3 = nn.Sequential(
nn.Conv2d(self.ngf * 2,
self.ngf * 4,
kernel_size=4,
stride=2,
padding=1,
bias=False), nn.InstanceNorm2d(self.ngf * 4,
affine=True),
nn.LeakyReLU(0.2))
self.down_4 = nn.Sequential(
nn.Conv2d(self.ngf * 4,
self.ngf * 8,
kernel_size=4,
stride=2,
padding=1,
bias=False), nn.InstanceNorm2d(self.ngf * 8,
affine=True),
nn.LeakyReLU(0.2))
self.down_5 = nn.Sequential(
nn.Conv2d(self.ngf * 8,
self.ngf * 8,
kernel_size=4,
stride=2,
padding=1,
bias=False), nn.InstanceNorm2d(self.ngf * 8,
affine=True),
nn.LeakyReLU(0.2))
self.down_6 = nn.Sequential(
nn.Conv2d(self.ngf * 8,
self.ngf * 8,
kernel_size=4,
stride=2,
padding=1,
bias=False), nn.InstanceNorm2d(self.ngf * 8,
affine=True),
nn.LeakyReLU(0.2))
self.down_7 = nn.Sequential(
nn.Conv2d(self.ngf * 8,
self.ngf * 8,
kernel_size=4,
stride=2,
padding=1,
bias=False), nn.InstanceNorm2d(self.ngf * 8,
affine=True),
nn.LeakyReLU(0.2))
self.up_1 = nn.Sequential(
nn.ConvTranspose2d(self.ngf * 8,
self.ngf * 8,
kernel_size=4,
stride=2,
padding=1,
bias=False),
nn.InstanceNorm2d(self.ngf * 8, affine=True), nn.ReLU())
self.up_2 = nn.Sequential(
nn.ConvTranspose2d(self.ngf * 16,
self.ngf * 8,
kernel_size=4,
stride=2,
padding=1,
bias=False),
nn.InstanceNorm2d(self.ngf * 8, affine=True), nn.ReLU())
self.up_3 = nn.Sequential(
nn.ConvTranspose2d(self.ngf * 16,
self.ngf * 8,
kernel_size=4,
stride=2,
padding=1,
bias=False),
nn.InstanceNorm2d(self.ngf * 8, affine=True), nn.ReLU())
self.up_4 = nn.Sequential(
nn.ConvTranspose2d(self.ngf * 16,
self.ngf * 4,
kernel_size=4,
stride=2,
padding=1,
bias=False),
nn.InstanceNorm2d(self.ngf * 4, affine=True), nn.ReLU())
self.up_5 = nn.Sequential(
nn.ConvTranspose2d(self.ngf * 8,
self.ngf * 2,
kernel_size=4,
stride=2,
padding=1,
bias=False),
nn.InstanceNorm2d(self.ngf * 2, affine=True), nn.ReLU())
self.up_6 = nn.Sequential(
nn.ConvTranspose2d(self.ngf * 4,
self.ngf,
kernel_size=4,
stride=2,
padding=1,
bias=False),
nn.InstanceNorm2d(self.ngf, affine=True), nn.ReLU())
self.up_7 = nn.Sequential(
nn.ConvTranspose2d(self.ngf * 2,
self.out_channels,
kernel_size=4,
stride=2,
padding=1,
bias=False), nn.Tanh())
def __init__(self,
input_dim,
output_dim,
kernel_size,
stride,
padding=0,
norm='none',
activation='relu',
pad_type='zero'):
super(Conv2dBlock, self).__init__()
self.use_bias = True
# initialize padding
if pad_type == 'reflect':
self.pad = nn.ReflectionPad2d(padding)
elif pad_type == 'replicate':
self.pad = nn.ReplicationPad2d(padding)
elif pad_type == 'zero':
self.pad = nn.ZeroPad2d(padding)
else:
assert 0, "Unsupported padding type: {}".format(pad_type)
# initialize normalization
norm_dim = output_dim
if norm == 'bn':
self.norm = nn.BatchNorm2d(norm_dim)
elif norm == 'in':
#self.norm = nn.InstanceNorm2d(norm_dim, track_running_stats=True)
self.norm = nn.InstanceNorm2d(norm_dim)
elif norm == 'ln':
self.norm = LayerNorm(norm_dim)
elif norm == 'adain':
self.norm = AdaptiveInstanceNorm2d(norm_dim)
elif norm == 'none' or norm == 'sn':
self.norm = None
else:
assert 0, "Unsupported normalization: {}".format(norm)
# initialize activation
if activation == 'relu':
self.activation = nn.ReLU(inplace=True)
elif activation == 'lrelu':
self.activation = nn.LeakyReLU(0.2, inplace=True)
elif activation == 'prelu':
self.activation = nn.PReLU()
elif activation == 'selu':
self.activation = nn.SELU(inplace=True)
elif activation == 'tanh':
self.activation = nn.Tanh()
elif activation == 'none':
self.activation = None
else:
assert 0, "Unsupported activation: {}".format(activation)
# initialize convolution
if norm == 'sn':
self.conv = SpectralNorm(
nn.Conv2d(input_dim,
output_dim,
kernel_size,
stride,
bias=self.use_bias))
else:
self.conv = nn.Conv2d(input_dim,
output_dim,
kernel_size,
stride,
bias=self.use_bias)
def __init__(self, n_classes, num_filter, BatchNorm, in_channels=3):
super(DilatedResnetForFlatByClassifyWithRgressV2v6v4c1GN, self).__init__()
self.in_channels = in_channels
self.n_classes = n_classes
self.num_filter = num_filter
# act_fn = nn.PReLU()
act_fn = nn.ReLU(inplace=True)
# act_fn = nn.LeakyReLU(0.2)
map_num = [1, 2, 4, 8, 16]
GN_num = self.num_filter * map_num[0]
is_dropout = False
print("\n------load DilatedResnetForFlatByClassifyWithRgressV2v6v4c1GN------\n")
self.resnet_head = nn.Sequential(
nn.Conv2d(self.in_channels, self.num_filter * map_num[0], kernel_size=7, stride=1, padding=3),
nn.InstanceNorm2d(self.num_filter * map_num[0]),
# BatchNorm(1, self.num_filter * map_num[0]),
# BatchNorm(self.num_filter * map_num[0]),
# nn.BatchNorm2d(self.num_filter * map_num[0]),
act_fn,
# nn.Dropout(p=0.2),
nn.Conv2d(self.num_filter * map_num[0], self.num_filter * map_num[0], kernel_size=7, stride=2, padding=3),
nn.InstanceNorm2d(self.num_filter * map_num[0]),
# BatchNorm(1, self.num_filter * map_num[0]),
# BatchNorm(self.num_filter * map_num[0]),
# nn.BatchNorm2d(self.num_filter * map_num[0]),
act_fn,
# nn.Dropout(p=0.2),
# nn.MaxPool2d(kernel_size=2, stride=2, padding=0),
nn.Conv2d(self.num_filter * map_num[0], self.num_filter * map_num[0], kernel_size=7, stride=2, padding=3),
nn.InstanceNorm2d(self.num_filter * map_num[0]),
# BatchNorm(1, self.num_filter * map_num[0]),
# BatchNorm(self.num_filter * map_num[0]),
# nn.BatchNorm2d(self.num_filter * map_num[0]),
act_fn,
# nn.Dropout(p=0.2),
)
# self.resnet_down = ResNetGN(num_filter, map_num, BatchNorm, GN_num=[32, 32, 32, 32], block_nums=[3, 4, 6, 3], block=ResidualBlock34GN, dropRate=[0, 0, 0, 0])
self.resnet_down = ResNetV2StraightV2GN(num_filter, map_num, BatchNorm, GN_num=[GN_num, GN_num, GN_num, GN_num], block_nums=[3, 4, 6, 3], block=ResidualBlock34DilatedV4GN, dropRate=[0, 0, 0, 0], is_sub_dropout=False)
map_num_i = 3
self.bridge_1 = nn.Sequential(
dilation_conv_gn_act(self.num_filter * map_num[map_num_i], self.num_filter * map_num[map_num_i],
act_fn, BatchNorm, GN_num=GN_num, dilation=1),
# conv_bn_act(self.num_filter * map_num[map_num_i], self.num_filter * map_num[map_num_i], act_fn),
)
self.bridge_2 = nn.Sequential(
dilation_conv_gn_act(self.num_filter * map_num[map_num_i], self.num_filter * map_num[map_num_i],
act_fn, BatchNorm, GN_num=GN_num, dilation=2),
# conv_bn_act(self.num_filter * map_num[map_num_i], self.num_filter * map_num[map_num_i], act_fn),
)
self.bridge_3 = nn.Sequential(
dilation_conv_gn_act(self.num_filter * map_num[map_num_i], self.num_filter * map_num[map_num_i],
act_fn, BatchNorm, GN_num=GN_num, dilation=5),
# conv_bn_act(self.num_filter * map_num[map_num_i], self.num_filter * map_num[map_num_i], act_fn),
)
self.bridge_4 = nn.Sequential(
dilation_conv_gn_act(self.num_filter * map_num[map_num_i], self.num_filter * map_num[map_num_i],
act_fn, BatchNorm, GN_num=GN_num, dilation=8),
dilation_conv_gn_act(self.num_filter * map_num[map_num_i], self.num_filter * map_num[map_num_i],
act_fn, BatchNorm, GN_num=GN_num, dilation=3),
dilation_conv_gn_act(self.num_filter * map_num[map_num_i], self.num_filter * map_num[map_num_i],
act_fn, BatchNorm, GN_num=GN_num, dilation=2),
)
self.bridge_5 = nn.Sequential(
dilation_conv_gn_act(self.num_filter * map_num[map_num_i], self.num_filter * map_num[map_num_i],
act_fn, BatchNorm, GN_num=GN_num, dilation=12),
dilation_conv_gn_act(self.num_filter * map_num[map_num_i], self.num_filter * map_num[map_num_i],
act_fn, BatchNorm, GN_num=GN_num, dilation=7),
dilation_conv_gn_act(self.num_filter * map_num[map_num_i], self.num_filter * map_num[map_num_i],
act_fn, BatchNorm, GN_num=GN_num, dilation=4),
)
self.bridge_6 = nn.Sequential(
dilation_conv_gn_act(self.num_filter * map_num[map_num_i], self.num_filter * map_num[map_num_i],
act_fn, BatchNorm, GN_num=GN_num, dilation=18),
dilation_conv_gn_act(self.num_filter * map_num[map_num_i], self.num_filter * map_num[map_num_i],
act_fn, BatchNorm, GN_num=GN_num, dilation=12),
dilation_conv_gn_act(self.num_filter * map_num[map_num_i], self.num_filter * map_num[map_num_i],
act_fn, BatchNorm, GN_num=GN_num, dilation=6),
)
self.bridge_concate = nn.Sequential(
nn.Conv2d(self.num_filter * map_num[map_num_i] * 6, self.num_filter * map_num[4], kernel_size=1, stride=1, padding=0),
BatchNorm(GN_num, self.num_filter * map_num[4]),
# BatchNorm(self.num_filter * map_num[4]),
# nn.BatchNorm2d(self.num_filter * map_num[4]),
act_fn,
)
self.regess_4 = ConvBlockResidualGN(self.num_filter * map_num[4], self.num_filter * (map_num[4]), act_fn, BatchNorm, GN_num=GN_num, is_dropout=False)
# self.regess_4 = ConvBlockV2(self.num_filter * map_num[4], self.num_filter * (map_num[3]), act_fn, BatchNorm, is_dropout=False)
# self.regess_4 = MergeBlockV2(self.num_filter * map_num[4], self.num_filter * (map_num[3]), n_classes, act_fn, BatchNorm, is_dropout=False)
# self.trans_4 = upsamplingBilinear(scale_factor=2)
self.trans_4 = transitionUpGN(self.num_filter * (map_num[4]), self.num_filter * map_num[3], act_fn, BatchNorm, GN_num=GN_num)
self.regess_3 = ConvBlockResidualGN(self.num_filter * (map_num[3]), self.num_filter * (map_num[3]), act_fn, BatchNorm, GN_num=GN_num, is_dropout=False)
# self.regess_3 = ConvBlockV2(self.num_filter * (map_num[3]), self.num_filter * (map_num[2]), act_fn, BatchNorm, is_dropout=False)
# self.regess_3 = MergeBlockV2(self.num_filter * (map_num[3]), self.num_filter * (map_num[2]), n_classes, act_fn, BatchNorm, is_dropout=False)
# self.trans_3 = upsamplingBilinear(scale_factor=2)
self.trans_3 = transitionUpGN(self.num_filter * (map_num[3]), self.num_filter * map_num[2], act_fn, BatchNorm, GN_num=GN_num)
self.regess_2 = ConvBlockResidualGN(self.num_filter * (map_num[2]), self.num_filter * (map_num[2]), act_fn, BatchNorm, GN_num=GN_num, is_dropout=False)
# self.regess_2 = ConvBlockV2(self.num_filter * (map_num[2]), self.num_filter * (map_num[2]), act_fn, BatchNorm, is_dropout=False)
# self.regess_2 = MergeBlockV2(self.num_filter * (map_num[2]), self.num_filter * (map_num[2]), n_classes, act_fn, BatchNorm, is_dropout=False)
# self.trans_2 = upsamplingBilinear(scale_factor=2)
self.trans_2 = transitionUpGN(self.num_filter * map_num[2], self.num_filter * map_num[1], act_fn, BatchNorm, GN_num=GN_num)
self.regess_1 = ConvBlockResidualGN(self.num_filter * (map_num[1]), self.num_filter * (map_num[1]), act_fn, BatchNorm, GN_num=GN_num, is_dropout=False)
# self.regess_1 = ConvBlockV2(self.num_filter * (map_num[2]), self.num_filter * (map_num[1]), act_fn, BatchNorm, is_dropout=False)
# self.regess_1 = MergeBlockV2(self.num_filter * (map_num[2]), self.num_filter * (map_num[1]), n_classes, act_fn, BatchNorm, is_dropout=False)
self.trans_1 = upsamplingBilinear(scale_factor=2)
# self.trans_1 = transitionUpGN(self.num_filter * map_num[1], self.num_filter * map_num[1], act_fn, BatchNorm, GN_num=GN_num)
self.regess_0 = ConvBlockResidualGN(self.num_filter * (map_num[1]), self.num_filter * (map_num[1]), act_fn, BatchNorm, GN_num=GN_num, is_dropout=False)
# self.regess_0 = ConvBlockV2(self.num_filter * (map_num[1]), self.num_filter * (map_num[0]), act_fn, BatchNorm, is_dropout=False)
# self.regess_0 = MergeBlockV2(self.num_filter * (map_num[1]), self.num_filter * (map_num[0]), n_classes, act_fn, BatchNorm, is_dropout=False)
# self.trans_0 = transitionUpGN(self.num_filter * map_num[1], self.num_filter * map_num[0], act_fn, BatchNorm, GN_num=GN_num)
self.trans_0 = upsamplingBilinear(scale_factor=2)
self.up = nn.Sequential(
nn.Conv2d(self.num_filter * map_num[1], self.num_filter * map_num[0], kernel_size=3, stride=1, padding=1),
# BatchNorm(GN_num, self.num_filter * map_num[0]),
nn.InstanceNorm2d(self.num_filter * map_num[0]),
# nn.BatchNorm2d(self.num_filter * map_num[0]),
act_fn,
# nn.Dropout2d(p=0.2),
)
self.out_regress = nn.Sequential(
# nn.Conv2d(map_num[0], map_num[0], kernel_size=3, stride=1, padding=1, bias=False),
# nn.InstanceNorm2d(map_num[0]),
# nn.PReLU(),
nn.Conv2d(self.num_filter * map_num[0], self.num_filter * map_num[0], kernel_size=3, stride=1, padding=1),
# BatchNorm(GN_num, self.num_filter * map_num[0]),
nn.InstanceNorm2d(self.num_filter * map_num[0]),
nn.PReLU(),
# nn.Dropout2d(p=0.2),
nn.Conv2d(self.num_filter * map_num[0], n_classes, kernel_size=3, stride=1, padding=1),
nn.InstanceNorm2d(n_classes),
# BatchNorm(1, n_classes),
# BatchNorm(n_classes),
nn.PReLU(),
nn.Conv2d(n_classes, n_classes, kernel_size=3, stride=1, padding=1),
#nn.Conv2d(map_num[0], n_classes, kernel_size=3, stride=1, padding=1, bias=False),
)
self.out_classify = nn.Sequential(
# nn.Conv2d(map_num[0], map_num[0], kernel_size=3, stride=1, padding=1, bias=False),
# nn.InstanceNorm2d(map_num[0]),
# act_fn,
# nn.Dropout(p=0.2),
nn.Conv2d(self.num_filter * map_num[0], self.num_filter * map_num[0], kernel_size=3, stride=1, padding=1),
# BatchNorm(GN_num, self.num_filter * map_num[0]),
nn.InstanceNorm2d(self.num_filter * map_num[0]),
act_fn,
nn.Dropout2d(p=0.2), # nn.Dropout(p=0.2),
nn.Conv2d(self.num_filter * map_num[0], n_classes, kernel_size=3, stride=1, padding=1),
nn.InstanceNorm2d(n_classes),
# BatchNorm(1, n_classes),
# BatchNorm(n_classes),
act_fn,
# nn.Dropout(p=0.2),
nn.Conv2d(n_classes, 1, kernel_size=3, stride=1, padding=1),
#nn.Conv2d(self.num_filter * map_num[0], n_classes, kernel_size=3, stride=1, padding=1, bias=False),
)
self.out_classify_softmax = nn.Sigmoid()
#self.out_classify_softmax = nn.Softmax2d()
self._initialize_weights()
def __init__(self, planes, ratio=0.5):
super(IBN, self).__init__()
self.half = int(planes * ratio)
self.IN = nn.InstanceNorm2d(self.half, affine=True)
self.BN = nn.BatchNorm2d(planes - self.half)
def __init__(self, input_channels, kernel_size, bias=True):
super(Conv2LSTM, self).__init__()
self.input_channels = input_channels
self.bias = bias
self.kernel_size = kernel_size
self.FME = nn.Sequential(
ConvScale(self.input_channels, 64, self.kernel_size, bias=True),
nn.InstanceNorm2d(64, affine=True),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=2,
stride=2,
padding=0,
dilation=1,
ceil_mode=False),
ConvScale(64, 128, self.kernel_size, bias=True),
nn.InstanceNorm2d(128, affine=True),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=2,
stride=2,
padding=0,
dilation=1,
ceil_mode=False),
ConvScale(128, 128, self.kernel_size, bias=True),
nn.InstanceNorm2d(128, affine=True),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=2,
stride=2,
padding=0,
dilation=1,
ceil_mode=False),
ConvScale(128, 64, self.kernel_size, bias=True),
nn.InstanceNorm2d(64, affine=True),
nn.ReLU(inplace=True),
#add
nn.Conv2d(64, 1, 1, stride=1, padding=int((1 - 1) / 2), bias=True),
nn.InstanceNorm2d(1, affine=True),
nn.ReLU(inplace=True),
)
self.DME = nn.Sequential(
nn.Conv2d(64, 64, 9, stride=1, padding=int((9 - 1) / 2),
bias=True),
nn.InstanceNorm2d(64, affine=True),
nn.ReLU(inplace=True),
nn.ConvTranspose2d(64,
64,
9,
stride=2,
padding=4,
output_padding=1,
bias=True),
nn.InstanceNorm2d(64, affine=True),
nn.ReLU(inplace=True),
nn.Conv2d(64, 32, 7, stride=1, padding=int((7 - 1) / 2),
bias=True),
nn.InstanceNorm2d(32, affine=True),
nn.ReLU(inplace=True),
nn.ConvTranspose2d(32,
32,
7,
stride=2,
padding=3,
output_padding=1,
bias=True),
nn.InstanceNorm2d(32, affine=True),
nn.ReLU(inplace=True),
nn.Conv2d(32, 16, 5, stride=1, padding=int((5 - 1) / 2),
bias=True),
nn.InstanceNorm2d(16, affine=True),
nn.ReLU(inplace=True),
nn.ConvTranspose2d(16,
16,
5,
stride=2,
padding=2,
output_padding=1,
bias=True),
nn.InstanceNorm2d(16, affine=True),
nn.ReLU(inplace=True),
nn.Conv2d(16, 16, 3, stride=1, padding=int((3 - 1) / 2),
bias=True),
nn.InstanceNorm2d(16, affine=True),
nn.ReLU(inplace=True),
nn.Conv2d(16, 16, 5, stride=1, padding=int((5 - 1) / 2),
bias=True),
nn.InstanceNorm2d(16, affine=True),
nn.ReLU(inplace=True),
nn.Conv2d(16, 1, 1, stride=1, padding=int((1 - 1) / 2), bias=True),
nn.InstanceNorm2d(1, affine=True),
nn.ReLU(inplace=True),
)
self.LSTM = nn.LSTM(input_size=int(480 * 640 / 64),
hidden_size=60 * 80,
num_layers=4,
batch_first=True,
dropout=0.5)
def __init__(self):
super(InstanceNorm, self).__init__()
self.inorm2 = nn.InstanceNorm2d(3, affine=True)
def __init__(
self,
input_dim,
output_dim,
kernel_size,
stride,
padding=0,
norm="none",
activation="relu",
pad_type="zero",
):
super(ConvBlock, self).__init__()
self.use_bias = True
# initialize padding
if pad_type == "reflect":
self.pad = nn.ReflectionPad2d(padding)
elif pad_type == "replicate":
self.pad = nn.ReplicationPad2d(padding)
elif pad_type == "zero":
self.pad = nn.ZeroPad2d(padding)
else:
assert 0, "Unsupported padding type: {}".format(pad_type)
# initialize normalization
norm_dim = output_dim
if norm == "bn":
self.norm = nn.BatchNorm2d(norm_dim)
elif norm == "in":
# self.norm = nn.InstanceNorm2d(norm_dim, track_running_stats=True)
self.norm = nn.InstanceNorm2d(norm_dim)
elif norm == "ln":
self.norm = LayerNorm(norm_dim)
elif norm == "adain":
self.norm = AdaptiveInstanceNorm2d(norm_dim)
elif norm == "none" or norm == "sn":
self.norm = None
else:
assert 0, "Unsupported normalization: {}".format(norm)
# initialize activation
if activation == "relu":
self.activation = nn.ReLU(inplace=True)
elif activation == "lrelu":
self.activation = nn.LeakyReLU(0.2, inplace=True)
elif activation == "prelu":
self.activation = nn.PReLU()
elif activation == "selu":
self.activation = nn.SELU(inplace=True)
elif activation == "tanh":
self.activation = nn.Tanh()
elif activation == "none":
self.activation = None
else:
assert 0, "Unsupported activation: {}".format(activation)
# initialize convolution
if norm == "sn":
self.conv = SpectralNorm(
nn.Conv2d(input_dim,
output_dim,
kernel_size,
stride,
bias=self.use_bias))
else:
self.conv = nn.Conv2d(input_dim,
output_dim,
kernel_size,
stride,
bias=self.use_bias)
def block(in_, out_, normalization=True):
layers = [nn.Conv2d(in_, out_, 4, stride=2, padding=1)]
if normalization:
layers.append(nn.InstanceNorm2d(out_))
layers.append(nn.LeakyReLU(0.2, inplace=True))
return layers
def get_normalization(name, channels, **kwargs):
if name == 'bn': return nn.BatchNorm2d(channels, **kwargs)
elif name == 'in': return nn.InstanceNorm2d(channels, **kwargs)
else: return nn.Identity()
def get_norm(num_features, norm_type):
if norm_type == 'BN':
return nn.BatchNorm2d(num_features)
elif norm_type == 'IN':
return nn.InstanceNorm2d(num_features)
def _make_stage(self, features, out_features, size):
prior = nn.AdaptiveAvgPool2d(output_size=(size, size))
conv = nn.Conv2d(features, out_features, kernel_size=1, bias=False)
bn = nn.InstanceNorm2d(out_features)
return nn.Sequential(prior, conv, bn)
def __init__(self, n_channel):
super().__init__()
self.norm = nn.InstanceNorm2d(n_channel)
def test_g():
padding = nn.ReflectionPad2d((3, 3, 3, 3))
relu = nn.ReLU()
conv1 = nn.Conv2d(in_channels=1,
out_channels=64,
kernel_size=7,
stride=1,
padding=0,
bias=False)
init.normal_(conv1.weight)
instance_norm = nn.InstanceNorm2d(64, eps=1e-5)
cnet1 = nn.Sequential(padding, conv1, instance_norm, relu)
conv2 = nn.Conv2d(in_channels=64,
out_channels=128,
kernel_size=3,
stride=2,
padding=1,
bias=False)
init.normal_(conv2.weight)
instance_norm = nn.InstanceNorm2d(128, eps=1e-5)
cnet2 = nn.Sequential(conv2, instance_norm, relu)
conv3 = nn.Conv2d(in_channels=128,
out_channels=256,
kernel_size=3,
stride=2,
padding=1,
bias=False)
init.normal_(conv3.weight)
instance_norm = nn.InstanceNorm2d(256, eps=1e-5)
cnet3 = nn.Sequential(conv3, instance_norm, relu)
resnet = nn.Sequential()
for i in range(10):
resnet.add_module('resnet_block', ResnetBlock(dim=256,
padding_type='reflect',
use_dropout=False,
use_bias=False,
norm_layer=nn.InstanceNorm2d))
extra_padding = nn.ZeroPad2d((0, 1, 0, 1))
tconv1 = nn.ConvTranspose2d(in_channels=256,
out_channels=128,
kernel_size=3,
stride=2,
padding=1,
bias=False)
init.normal_(tconv1.weight)
tcnet1 = nn.Sequential(tconv1, extra_padding, instance_norm, relu)
tconv2 = nn.ConvTranspose2d(in_channels=128,
out_channels=64,
kernel_size=3,
stride=2,
padding=1,
bias=False)
init.normal_(tconv2.weight)
tcnet2 = nn.Sequential(tconv2, extra_padding, instance_norm, relu)
conv4 = nn.Conv2d(in_channels=64,
out_channels=1,
kernel_size=7,
stride=1,
padding=0)
init.normal_(conv4.weight)
sigmoid = nn.Sigmoid()
cnet4 = nn.Sequential(padding, conv4, sigmoid)
x = torch.ones((4, 1, 64, 84))
x = cnet1(x)
x = cnet2(x)
x = cnet3(x)
x = resnet(x)
x = tcnet1(x)
x = tcnet2(x)
x = cnet4(x)
print(x.shape)
def __init__(self):
super(Discriminator, self).__init__()
self.in_channels = 3
self.ndf = 64
self.out_channels = 1
# Discriminator having last patch of (1, 13, 13)
# (batch_size, 3, 128, 128) -> (batch_size, 1, 13, 13)
self.main_1 = nn.Sequential(
nn.AvgPool2d(kernel_size=3,
stride=2,
padding=0,
count_include_pad=False),
nn.Conv2d(self.in_channels,
32,
kernel_size=4,
stride=2,
padding=1,
bias=False), nn.LeakyReLU(0.2),
nn.Conv2d(32, 64, kernel_size=4, stride=2, padding=1, bias=False),
nn.InstanceNorm2d(64, affine=True), nn.LeakyReLU(0.2),
nn.Conv2d(64, 128, kernel_size=4, stride=1, padding=1, bias=False),
nn.InstanceNorm2d(128, affine=True), nn.LeakyReLU(0.2),
nn.Conv2d(128,
self.out_channels,
kernel_size=4,
stride=1,
padding=1,
bias=False))
# Discriminator having last patch of (1, 30, 30)
# (batch_size, 3, 128, 128) -> (batch_size, 1, 30, 30)
self.main_2 = nn.Sequential(
nn.Conv2d(self.in_channels,
self.ndf,
kernel_size=4,
stride=2,
padding=1,
bias=False), nn.LeakyReLU(0.2),
nn.Conv2d(self.ndf,
self.ndf * 2,
kernel_size=4,
stride=2,
padding=1,
bias=False), nn.InstanceNorm2d(self.ndf * 2,
affine=True),
nn.LeakyReLU(0.2),
nn.Conv2d(self.ndf * 2,
self.ndf * 4,
kernel_size=4,
stride=1,
padding=1,
bias=False), nn.InstanceNorm2d(self.ndf * 4,
affine=True),
nn.LeakyReLU(0.2),
nn.Conv2d(self.ndf * 4,
self.out_channels,
kernel_size=4,
stride=1,
padding=1,
bias=False))
def __init__(self):
super(Generator, self).__init__()
self.net1 = nn.Sequential(nn.ReflectionPad2d((3, 3, 3, 3)),
nn.Conv2d(in_channels=1,
out_channels=64,
kernel_size=7,
stride=1,
padding=0,
bias=False),
nn.InstanceNorm2d(64, eps=1e-5),
nn.ReLU()
)
self.net2 = nn.Sequential(nn.Conv2d(in_channels=64,
out_channels=128,
kernel_size=3,
stride=2,
padding=1,
bias=False),
nn.InstanceNorm2d(128, eps=1e-5),
nn.ReLU(),
nn.Conv2d(in_channels=128,
out_channels=256,
kernel_size=3,
stride=2,
padding=1,
bias=False),
nn.InstanceNorm2d(256, eps=1e-5),
nn.ReLU()
)
self.resnet = nn.Sequential()
for i in range(10):
self.resnet.add_module('resnet_block', ResnetBlock(dim=256,
padding_type='reflect',
use_dropout=False,
use_bias=False,
norm_layer=nn.InstanceNorm2d))
self.net3 = nn.Sequential(nn.ConvTranspose2d(in_channels=256,
out_channels=128,
kernel_size=3,
stride=2,
padding=1,
bias=False),
nn.ZeroPad2d((0, 1, 0, 1)),
nn.InstanceNorm2d(128, eps=1e-5),
nn.ReLU(),
nn.ConvTranspose2d(in_channels=128,
out_channels=64,
kernel_size=3,
stride=2,
padding=1,
bias=False),
nn.ZeroPad2d((0, 1, 0, 1)),
nn.InstanceNorm2d(64, eps=1e-5),
nn.ReLU(),
nn.ReflectionPad2d((3, 3, 3, 3)),
nn.Conv2d(in_channels=64,
out_channels=1,
kernel_size=7,
stride=1,
padding=0),
nn.Sigmoid()
)
def __init__(self, norm_nc, cond_nc):
super().__init__()
self.norm = nn.InstanceNorm2d(norm_nc, affine=False)
def __init__(self):
super(GeneratorStructural, self).__init__()
self.fc1 = nn.Linear(128, 4 * 4 * 1024)
self.fc1_bn = nn.InstanceNorm2d(1024)
self.res_block1 = ResidualBlock('G', 1024, 512, resample='up')
self.relu = nn.ReLU()
def test_instance_norm(self):
underlying = nn.InstanceNorm2d(3)
self.run_model_test(underlying, train=False, batch_size=BATCH_SIZE)
def __init__(self, name, in_channels):
super(Normalize, self).__init__()
if name == 'D':
self.norm = Identity()
else:
self.norm = nn.InstanceNorm2d(in_channels, affine=True)
def __init__(
self,
in_channels,
out_channels,
kernel_size,
stride=1,
padding=0,
dilation=1,
pad_type="replicate",
activation="elu",
norm="none",
sc=False,
):
super().__init__()
# Initialize the padding scheme
if pad_type == "reflect":
self.pad = nn.ReflectionPad2d(padding)
elif pad_type == "replicate":
self.pad = nn.ReplicationPad2d(padding)
elif pad_type == "zero":
self.pad = nn.ZeroPad2d(padding)
else:
assert 0, f"Unsupported padding type: {pad_type}"
# Initialize the normalization type
if norm == "bn":
self.norm = nn.BatchNorm2d(out_channels)
elif norm == "in":
self.norm = nn.InstanceNorm2d(out_channels)
elif norm == "none":
self.norm = None
else:
assert 0, f"Unsupported normalization: {norm}"
# Initialize the activation funtion
if activation == "relu":
self.activation = nn.ReLU(inplace=True)
elif activation == "elu":
self.activation = nn.ELU(alpha=1.0, inplace=True)
elif activation == "lrelu":
self.activation = nn.LeakyReLU(0.2, inplace=True)
elif activation == "prelu":
self.activation = nn.PReLU()
elif activation == "selu":
self.activation = nn.SELU(inplace=True)
elif activation == "tanh":
self.activation = nn.Tanh()
elif activation == "sigmoid":
self.activation = nn.Sigmoid()
elif activation == "none":
self.activation = None
else:
assert 0, f"Unsupported activation: {activation}"
# Initialize the convolution layers
if sc:
self.conv2d = nn.Conv2d(in_channels,
out_channels,
kernel_size,
stride,
padding=0,
dilation=dilation)
self.mask_conv2d = sc_conv(in_channels,
kernel_size,
stride,
padding=0,
dilation=dilation)
else:
self.conv2d = nn.Conv2d(in_channels,
out_channels,
kernel_size,
stride,
padding=0,
dilation=dilation)
# self.mask_conv2d = nn.Conv2d(in_channels, out_channels, kernel_size, stride,
# padding=0, dilation=dilation)
self.mask_conv2d = depth_separable_conv(in_channels,
out_channels,
kernel_size,
stride,
padding=0,
dilation=dilation)
self.sigmoid = torch.nn.Sigmoid()
def __init__(self,
n_intervals,
n_blocks=16,
inchannels=3,
nfeats=64,
outchannels=3):
super(_NetG_SAM, self).__init__()
self.n_blocks = n_blocks
self.intervals = n_intervals
if isinstance(n_intervals, list):
self.nbody = len(n_intervals)
if isinstance(n_intervals, int):
self.nbody = self.n_blocks // n_intervals
self.conv_input = nn.Conv2d(in_channels=3,
out_channels=64,
kernel_size=9,
stride=1,
padding=4,
bias=False)
self.relu = nn.LeakyReLU(0.2, inplace=True)
self.residual = self.make_layer(_Residual_Block, 16)
self.conv_mid = nn.Conv2d(in_channels=64,
out_channels=64,
kernel_size=3,
stride=1,
padding=1,
bias=False)
self.bn_mid = nn.InstanceNorm2d(64, affine=True)
self.upscale4x = nn.Sequential(
nn.Conv2d(in_channels=64,
out_channels=256,
kernel_size=3,
stride=1,
padding=1,
bias=False),
nn.PixelShuffle(2),
nn.LeakyReLU(0.2, inplace=True),
nn.Conv2d(in_channels=64,
out_channels=256,
kernel_size=3,
stride=1,
padding=1,
bias=False),
nn.PixelShuffle(2),
nn.LeakyReLU(0.2, inplace=True),
)
self.conv_output = nn.Conv2d(in_channels=64,
out_channels=3,
kernel_size=9,
stride=1,
padding=4,
bias=False)
sam_layer = []
for _ in range(self.nbody):
sam_layer.append(SAM(nfeats))
self.sam_layer = nn.Sequential(*sam_layer)
def __init__(self,
num_classes,
blocks,
layers,
channels,
feature_dim=256,
loss='softmax',
instance_norm=False,
dropout_cfg=None,
fpn_cfg=None,
pooling_type='avg',
input_size=(256, 128),
IN_first=False,
extra_blocks=False,
lct_gate=False,
**kwargs):
self.dropout_cfg = dropout_cfg
self.extra_blocks = extra_blocks
self.channel_gate = LCTGate if lct_gate else ChannelGate
if self.extra_blocks:
for i, l in enumerate(layers):
layers[i] = l + 1
super(OSNetFPN, self).__init__(num_classes, blocks, layers, channels,
feature_dim, loss, instance_norm)
self.feature_scales = (4, 8, 16, 16)
if fpn_cfg is not None:
self.fpn_enable = fpn_cfg.enable
self.fpn_dim = fpn_cfg.dim
self.fpn_process = fpn_cfg.process
assert self.fpn_process in [
'concatenation', 'max_pooling', 'elementwise_sum'
]
else:
self.fpn_enable = False
self.feature_dim = feature_dim
self.use_IN_first = IN_first
if IN_first:
self.in_first = nn.InstanceNorm2d(3, affine=True)
self.conv1 = ConvLayer(3,
channels[0],
7,
stride=2,
padding=3,
IN=self.use_IN_first)
if self.fpn_enable:
self.fpn = FPN(channels, self.feature_scales, self.fpn_dim,
self.fpn_dim)
fpn_out_dim = self.fpn_dim if self.fpn_process in ['max_pooling', 'elementwise_sum'] \
else feature_dim
self.fc = self._construct_fc_layer(feature_dim, fpn_out_dim,
dropout_cfg)
else:
self.fpn = None
self.fc = self._construct_fc_layer(feature_dim, channels[3],
dropout_cfg)
if self.loss not in [
'am_softmax',
]:
self.classifier = nn.Linear(self.feature_dim, num_classes)
else:
self.classifier = AngleSimpleLinear(self.feature_dim, num_classes)
if 'conv' in pooling_type:
kernel_size = (input_size[0] // self.feature_scales[-1],
input_size[1] // self.feature_scales[-1])
if self.fpn_enable:
self.global_avgpool = nn.Conv2d(fpn_out_dim,
fpn_out_dim,
kernel_size,
groups=fpn_out_dim)
else:
self.global_avgpool = nn.Conv2d(channels[3],
channels[3],
kernel_size,
groups=channels[3])
elif 'avg' in pooling_type:
self.global_avgpool = nn.AdaptiveAvgPool2d(1)
elif 'gmp' in pooling_type:
self.global_avgpool = GeneralizedMeanPooling()
else:
raise ValueError('Incorrect pooling type')
if self.fpn_enable and self.fpn_process == 'concatenation':
self.fpn_extra_conv = ConvLayer(self.fpn_dim *
len(self.fpn.dims_out),
feature_dim,
3,
stride=1,
padding=1,
IN=False)
else:
self.fpn_extra_conv = None
self._init_params()
def __init__(self, num_features, norm_layer='in', eps=1e-4):
super(AdaptiveNorm2dTrainable, self).__init__()
self.num_features = num_features
if 'in' in norm_layer:
self.norm_layer = nn.InstanceNorm2d(num_features, eps=eps, affine=False)
def __init__(
self,
in_channels,
out_channels,
*,
bottleneck_channels,
stride=1,
num_groups=1,
norm="BN",
stride_in_1x1=False,
dilation=1,
instance_norm=False,
):
"""
Args:
bottleneck_channels (int): number of output channels for the 3x3
"bottleneck" conv layers.
num_groups (int): number of groups for the 3x3 conv layer.
norm (str or callable): normalization for all conv layers.
See :func:`layers.get_norm` for supported format.
stride_in_1x1 (bool): when stride>1, whether to put stride in the
first 1x1 convolution or the bottleneck 3x3 convolution.
dilation (int): the dilation rate of the 3x3 conv layer.
IN (bool): Flag for Instance Normalisation
"""
super().__init__(in_channels, out_channels, stride)
if in_channels != out_channels:
self.shortcut = Conv2d(
in_channels,
out_channels,
kernel_size=1,
stride=stride,
bias=False,
norm=get_norm(norm, out_channels),
)
else:
self.shortcut = None
# The original MSRA ResNet models have stride in the first 1x1 conv
# The subsequent fb.torch.resnet and Caffe2 ResNe[X]t implementations have
# stride in the 3x3 conv
stride_1x1, stride_3x3 = (stride, 1) if stride_in_1x1 else (1, stride)
self.conv1 = Conv2d(
in_channels,
bottleneck_channels,
kernel_size=1,
stride=stride_1x1,
bias=False,
norm=get_norm(norm, bottleneck_channels),
)
self.conv2 = Conv2d(
bottleneck_channels,
bottleneck_channels,
kernel_size=3,
stride=stride_3x3,
padding=1 * dilation,
bias=False,
groups=num_groups,
dilation=dilation,
norm=get_norm(norm, bottleneck_channels),
)
self.conv3 = Conv2d(
bottleneck_channels,
out_channels,
kernel_size=1,
bias=False,
norm=get_norm(norm, out_channels),
)
for layer in [self.conv1, self.conv2, self.conv3, self.shortcut]:
if layer is not None: # shortcut can be None
weight_init.c2_msra_fill(layer)
if instance_norm:
self.instance_norm = nn.InstanceNorm2d(out_channels, affine=True)
else:
self.instance_norm = None
def __init__(self, out_channels, track_running_stats=True):
super().__init__()
self.layer = nn.InstanceNorm2d(
out_channels, track_running_stats=track_running_stats)
