def __init__(self, configer):
super(asymmetric_non_local_network, self).__init__()
self.configer = configer
self.num_classes = self.configer.get('data', 'num_classes')
self.backbone = ModuleHelper.get_backbone(
backbone=self.configer.get('network.backbone'),
pretrained=self.configer.get('network.pretrained')
)
# low_in_channels, high_in_channels, out_channels, key_channels, value_channels, dropout
self.fusion = AFNB(1024, 2048, 2048, 256, 256, dropout=0.05, sizes=([1]), norm_type=self.configer.get('network', 'norm_type'))
# extra added layers
self.context = nn.Sequential(
nn.Conv2d(2048, 512, kernel_size=3, stride=1, padding=1),
ModuleHelper.BNReLU(512, norm_type=self.configer.get('network', 'norm_type')),
APNB(in_channels=512, out_channels=512, key_channels=256, value_channels=256,
dropout=0.05, sizes=([1]), norm_type=self.configer.get('network', 'norm_type'))
)
self.cls = nn.Conv2d(512, self.num_classes, kernel_size=1, stride=1, padding=0, bias=True)
self.dsn = nn.Sequential(
nn.Conv2d(1024, 512, kernel_size=3, stride=1, padding=1),
ModuleHelper.BNReLU(512, norm_type=self.configer.get('network', 'norm_type')),
nn.Dropout2d(0.05),
nn.Conv2d(512, self.num_classes, kernel_size=1, stride=1, padding=0, bias=True)
)
self.valid_loss_dict = configer.get('loss', 'loss_weights', configer.get('loss.loss_type'))
def __init__(self, input_num, num1, num2, dilation_rate, drop_out,
norm_type):
super(_DenseAsppBlock, self).__init__()
self.add_module(
'conv1',
nn.Conv2d(in_channels=input_num, out_channels=num1,
kernel_size=1)),
self.add_module(
'norm1',
ModuleHelper.BatchNorm2d(norm_type=norm_type)(num_features=num1)),
self.add_module('relu1', nn.ReLU(inplace=False)),
self.add_module(
'conv2',
nn.Conv2d(in_channels=num1,
out_channels=num2,
kernel_size=3,
dilation=dilation_rate,
padding=dilation_rate)),
self.add_module(
'norm2',
ModuleHelper.BatchNorm2d(norm_type=norm_type)(
num_features=input_num)),
self.add_module('relu2', nn.ReLU(inplace=False)),
self.drop_rate = drop_out
def __init__(self, low_in_channels, high_in_channels, key_channels, value_channels, out_channels=None, scale=1, norm_type=None,psp_size=(1,3,6,8)):
super(_SelfAttentionBlock, self).__init__()
self.scale = scale
self.in_channels = low_in_channels
self.out_channels = out_channels
self.key_channels = key_channels
self.value_channels = value_channels
if out_channels == None:
self.out_channels = high_in_channels
self.pool = nn.MaxPool2d(kernel_size=(scale, scale))
self.f_key = nn.Sequential(
nn.Conv2d(in_channels=self.in_channels, out_channels=self.key_channels,
kernel_size=1, stride=1, padding=0),
ModuleHelper.BNReLU(self.key_channels, norm_type=norm_type),
)
self.f_query = nn.Sequential(
nn.Conv2d(in_channels=high_in_channels, out_channels=self.key_channels,
kernel_size=1, stride=1, padding=0),
ModuleHelper.BNReLU(self.key_channels, norm_type=norm_type),
)
self.f_value = nn.Conv2d(in_channels=self.in_channels, out_channels=self.value_channels,
kernel_size=1, stride=1, padding=0)
self.W = nn.Conv2d(in_channels=self.value_channels, out_channels=self.out_channels,
kernel_size=1, stride=1, padding=0)
self.psp = PSPModule(psp_size)
nn.init.constant_(self.W.weight, 0)
nn.init.constant_(self.W.bias, 0)
def __init__(self, input_nc, ndf=64, n_layers=3, norm_type=None):
"""Construct a PatchGAN discriminator
Parameters:
input_nc (int) -- the number of channels in input images
ndf (int) -- the number of filters in the last conv layer
n_layers (int) -- the number of conv layers in the discriminator
norm_layer -- normalization layer
"""
super(NLayerDiscriminator, self).__init__()
use_bias = (norm_type == 'instancenorm')
kw = 4
padw = 1
sequence = [nn.Conv2d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw), nn.LeakyReLU(0.2, True)]
nf_mult = 1
nf_mult_prev = 1
for n in range(1, n_layers): # gradually increase the number of filters
nf_mult_prev = nf_mult
nf_mult = min(2 ** n, 8)
sequence += [
nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=2, padding=padw, bias=use_bias),
ModuleHelper.BatchNorm2d(norm_type=norm_type)(ndf * nf_mult),
nn.LeakyReLU(0.2, True)
]
nf_mult_prev = nf_mult
nf_mult = min(2 ** n_layers, 8)
sequence += [
nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=1, padding=padw, bias=use_bias),
ModuleHelper.BatchNorm2d(norm_type=norm_type)(ndf * nf_mult),
nn.LeakyReLU(0.2, True)
]
sequence += [
nn.Conv2d(ndf * nf_mult, 1, kernel_size=kw, stride=1, padding=padw)] # output 1 channel prediction map
self.model = nn.Sequential(*sequence)
def __init__(self, configer):
super(ResSFNet, self).__init__()
self.configer = configer
self.num_classes = self.configer.get('data', 'num_classes')
base = ModuleHelper.get_backbone(
backbone=self.configer.get('network.backbone'),
pretrained=self.configer.get('network.pretrained'))
self.stage1 = nn.Sequential(base.conv1, base.bn1, base.relu1,
base.conv2, base.bn2, base.relu2,
base.conv3, base.bn3, base.relu3,
base.maxpool, base.layer1)
self.stage2 = base.layer2
self.stage3 = base.layer3
self.stage4 = base.layer4
num_features = 512 if 'resnet18' in self.configer.get(
'network.backbone') else 2048
fpn_dim = max(num_features // 8, 128)
self.head = AlignHead(num_features, fpn_dim=fpn_dim)
self.dsn = nn.Sequential(
nn.Conv2d(num_features // 2,
max(num_features // 4, 256),
kernel_size=3,
stride=1,
padding=1),
ModuleHelper.BNReLU(max(num_features // 4, 256),
norm_type="batchnorm"), nn.Dropout2d(0.1),
nn.Conv2d(max(num_features // 4, 256),
self.num_classes,
kernel_size=1,
stride=1,
padding=0,
bias=True))
self.conv_last = nn.Sequential(
conv3x3_bn_relu(4 * fpn_dim, fpn_dim, 1),
nn.Conv2d(fpn_dim, self.num_classes, kernel_size=1))
self.fpn_dsn = nn.ModuleList()
for i in range(len([2, 4, 8])):
self.fpn_dsn.append(
nn.Sequential(
nn.Conv2d(fpn_dim,
fpn_dim,
kernel_size=3,
stride=1,
padding=1),
ModuleHelper.BNReLU(fpn_dim, norm_type="batchnorm"),
nn.Dropout2d(0.1),
nn.Conv2d(fpn_dim,
self.num_classes,
kernel_size=1,
stride=1,
padding=0,
bias=True)))
self.valid_loss_dict = configer.get('loss', 'loss_weights',
configer.get('loss.loss_type'))
def __init__(self,
low_in_channels,
high_in_channels,
out_channels,
key_channels,
value_channels,
dropout,
sizes=([1]),
norm_type=None,
psp_size=(1, 3, 6, 8)):
super(AFNB, self).__init__()
self.stages = []
self.norm_type = norm_type
self.psp_size = psp_size
self.stages = nn.ModuleList([
self._make_stage([low_in_channels, high_in_channels], out_channels,
key_channels, value_channels, size)
for size in sizes
])
self.conv_bn_dropout = nn.Sequential(
nn.Conv2d(out_channels + high_in_channels,
out_channels,
kernel_size=1,
padding=0),
ModuleHelper.BatchNorm2d(norm_type=self.norm_type)(out_channels),
nn.Dropout2d(dropout))
def build_conv_block(self, dim, padding_type, norm_type, use_dropout,
use_bias):
"""Construct a convolutional block.
Parameters:
dim (int) -- the number of channels in the conv layer.
padding_type (str) -- the name of padding layer: reflect | replicate | zero
norm_layer -- normalization layer
use_dropout (bool) -- if use dropout layers.
use_bias (bool) -- if the conv layer uses bias or not
Returns a conv block (with a conv layer, a normalization layer, and a non-linearity layer (ReLU))
"""
conv_block = []
p = 0
if padding_type == 'reflect':
conv_block += [nn.ReflectionPad2d(1)]
elif padding_type == 'replicate':
conv_block += [nn.ReplicationPad2d(1)]
elif padding_type == 'zero':
p = 1
else:
raise NotImplementedError('padding [%s] is not implemented' %
padding_type)
conv_block += [
nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias),
ModuleHelper.BatchNorm2d(norm_type=norm_type)(dim),
nn.ReLU(True)
]
if use_dropout:
conv_block += [nn.Dropout(0.5)]
p = 0
if padding_type == 'reflect':
conv_block += [nn.ReflectionPad2d(1)]
elif padding_type == 'replicate':
conv_block += [nn.ReplicationPad2d(1)]
elif padding_type == 'zero':
p = 1
else:
raise NotImplementedError('padding [%s] is not implemented' %
padding_type)
conv_block += [
nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias),
ModuleHelper.BatchNorm2d(norm_type=norm_type)(dim)
]
return nn.Sequential(*conv_block)
def __init__(self, configer):
super(BaseModel, self).__init__()
self.configer = configer
self.net = ModuleHelper.get_backbone(
backbone=configer.get('network.{}backbone'.format(self.flag)),
pretrained=configer.get('network.{}pretrained'.format(self.flag)))
self.valid_loss_dict = configer.get('loss.loss_weights',
configer.get('loss.loss_type'))
def __init__(self, configer):
super(OpenPose, self).__init__()
self.configer = configer
self.backbone = ModuleHelper.get_backbone(
backbone=self.configer.get('network.backbone'),
pretrained=self.configer.get('network.pretrained'))
self.pose_model = PoseModel(configer, self.backbone.get_num_features())
self.valid_loss_dict = configer.get('loss', 'loss_weights',
configer.get('loss.loss_type'))
def __init__(self, configer):
super(DarkNetYolov3, self).__init__()
self.configer = configer
self.backbone = ModuleHelper.get_backbone(
backbone=configer.get('network.backbone'),
pretrained=configer.get('network.pretrained', default=None))
self.yolov3_head = Yolov3Head(configer,
out_filters=self.backbone.num_features)
self.yolo_detection_layer = YOLODetectionLayer(self.configer)
def conv3x3_bn_relu(in_planes, out_planes, stride=1, norm_type="batchnorm"):
return nn.Sequential(
nn.Conv2d(in_planes,
out_planes,
kernel_size=3,
stride=stride,
padding=1,
bias=False),
ModuleHelper.BNReLU(out_planes, norm_type=norm_type),
)
def __init__(self, configer):
super(PSPNet, self).__init__()
self.configer = configer
self.num_classes = self.configer.get('data', 'num_classes')
base = ModuleHelper.get_backbone(
backbone=self.configer.get('network.backbone'),
pretrained=self.configer.get('network.pretrained'))
self.stage1 = nn.Sequential(base.conv1, base.bn1, base.relu1,
base.conv2, base.bn2, base.relu2,
base.conv3, base.bn3, base.relu3,
base.maxpool, base.layer1, base.layer2,
base.layer3)
self.stage2 = base.layer4
num_features = 512 if 'resnet18' in self.configer.get(
'network.backbone') else 2048
self.dsn = nn.Sequential(
_ConvBatchNormReluBlock(num_features // 2,
num_features // 4,
3,
1,
norm_type=self.configer.get(
'network', 'norm_type')),
nn.Dropout2d(0.1),
nn.Conv2d(num_features // 4, self.num_classes, 1, 1, 0))
self.ppm = PPMBilinearDeepsup(fc_dim=num_features,
norm_type=self.configer.get(
'network', 'norm_type'))
self.cls = nn.Sequential(
nn.Conv2d(num_features + 4 * 512,
512,
kernel_size=3,
padding=1,
bias=False),
ModuleHelper.BNReLU(512,
norm_type=self.configer.get(
'network', 'norm_type')),
nn.Dropout2d(0.1), nn.Conv2d(512, self.num_classes, kernel_size=1))
self.valid_loss_dict = configer.get('loss', 'loss_weights',
configer.get('loss.loss_type'))
def __init__(self, features, inner_features=512, out_features=512, dilations=(12, 24, 36), norm_type=None):
super(ASPPModule, self).__init__()
self.conv1 = nn.Sequential(nn.AdaptiveAvgPool2d((1, 1)),
nn.Conv2d(features, inner_features, kernel_size=1, padding=0, dilation=1,
bias=False),
ModuleHelper.BNReLU(inner_features, norm_type=norm_type))
self.conv2 = nn.Sequential(
nn.Conv2d(features, inner_features, kernel_size=1, padding=0, dilation=1, bias=False),
ModuleHelper.BNReLU(inner_features, norm_type=norm_type))
self.conv3 = nn.Sequential(
nn.Conv2d(features, inner_features, kernel_size=3, padding=dilations[0], dilation=dilations[0], bias=False),
ModuleHelper.BNReLU(inner_features, norm_type=norm_type))
self.conv4 = nn.Sequential(
nn.Conv2d(features, inner_features, kernel_size=3, padding=dilations[1], dilation=dilations[1], bias=False),
ModuleHelper.BNReLU(inner_features, norm_type=norm_type))
self.conv5 = nn.Sequential(
nn.Conv2d(features, inner_features, kernel_size=3, padding=dilations[2], dilation=dilations[2], bias=False),
ModuleHelper.BNReLU(inner_features, norm_type=norm_type))
self.bottleneck = nn.Sequential(
nn.Conv2d(inner_features * 5, out_features, kernel_size=1, padding=0, dilation=1, bias=False),
ModuleHelper.BNReLU(out_features, norm_type=norm_type),
nn.Dropout2d(0.1)
)
def __init__(self, configer):
self.inplanes = 128
super(DeepLabV3, self).__init__()
self.configer = configer
self.num_classes = self.configer.get('data', 'num_classes')
base = ModuleHelper.get_backbone(
backbone=self.configer.get('network.backbone'),
pretrained=self.configer.get('network.pretrained'))
self.stage1 = nn.Sequential(base.conv1, base.bn1, base.relu1,
base.conv2, base.bn2, base.relu2,
base.conv3, base.bn3, base.relu3,
base.maxpool, base.layer1, base.layer2,
base.layer3)
self.stage2 = base.layer4
num_features = 512 if 'resnet18' in self.configer.get(
'network.backbone') else 2048
self.head = nn.Sequential(
ASPPModule(num_features,
norm_type=self.configer.get('network', 'norm_type')),
nn.Conv2d(512,
self.num_classes,
kernel_size=1,
stride=1,
padding=0,
bias=True))
self.dsn = nn.Sequential(
nn.Conv2d(1024, 512, kernel_size=3, stride=1, padding=1),
ModuleHelper.BNReLU(512,
norm_type=self.configer.get(
'network', 'norm_type')),
nn.Dropout2d(0.1),
nn.Conv2d(512,
self.num_classes,
kernel_size=1,
stride=1,
padding=0,
bias=True))
self.valid_loss_dict = configer.get('loss', 'loss_weights',
configer.get('loss.loss_type'))
def __init__(self, num_input_features, num_output_features, norm_type):
super(_Transition, self).__init__()
self.add_module(
'conv',
nn.Conv2d(num_input_features,
num_output_features,
kernel_size=1,
stride=1,
bias=False))
self.add_module(
'norm',
ModuleHelper.BatchNorm2d(norm_type=norm_type)(
num_features=num_output_features)),
self.add_module('relu', nn.ReLU(inplace=False))
def __init__(self, configer):
super(DenseASPP, self).__init__()
self.configer = configer
dropout0 = MODEL_CONFIG['dropout0']
dropout1 = MODEL_CONFIG['dropout1']
d_feature0 = MODEL_CONFIG['d_feature0']
d_feature1 = MODEL_CONFIG['d_feature1']
self.backbone = ModuleHelper.get_backbone(
backbone=self.configer.get('network.backbone'),
pretrained=self.configer.get('network.pretrained')
)
num_features = self.backbone.get_num_features()
self.trans = _Transition(num_input_features=self.num_features,
num_output_features=self.num_features // 2,
norm_type=self.configer.get('network', 'norm_type'))
self.num_features = self.num_features // 2
self.ASPP_3 = _DenseAsppBlock(input_num=num_features, num1=d_feature0, num2=d_feature1,
dilation_rate=3, drop_out=dropout0,
norm_type=self.configer.get('network', 'norm_type'))
self.ASPP_6 = _DenseAsppBlock(input_num=num_features + d_feature1 * 1, num1=d_feature0, num2=d_feature1,
dilation_rate=6, drop_out=dropout0,
norm_type=self.configer.get('network', 'norm_type'))
self.ASPP_12 = _DenseAsppBlock(input_num=num_features + d_feature1 * 2, num1=d_feature0, num2=d_feature1,
dilation_rate=12, drop_out=dropout0,
norm_type=self.configer.get('network', 'norm_type'))
self.ASPP_18 = _DenseAsppBlock(input_num=num_features + d_feature1 * 3, num1=d_feature0, num2=d_feature1,
dilation_rate=18, drop_out=dropout0,
norm_type=self.configer.get('network', 'norm_type'))
self.ASPP_24 = _DenseAsppBlock(input_num=num_features + d_feature1 * 4, num1=d_feature0, num2=d_feature1,
dilation_rate=24, drop_out=dropout0,
norm_type=self.configer.get('network', 'norm_type'))
num_features = num_features + 5 * d_feature1
self.classification = nn.Sequential(
nn.Dropout2d(p=dropout1),
nn.Conv2d(num_features, self.configer.get('data', 'num_classes'), kernel_size=1, padding=0)
)
self.valid_loss_dict = configer.get('loss', 'loss_weights', configer.get('loss.loss_type'))
def __init__(self, fc_dim=4096, norm_type=None):
super(PPMBilinearDeepsup, self).__init__()
self.norm_type = norm_type
pool_scales = (1, 2, 3, 6)
self.ppm = []
# assert norm_type == 'syncbn' or not self.training
# Torch BN can't handle feature map size with 1x1.
for scale in pool_scales:
self.ppm.append(
nn.Sequential(
nn.AdaptiveAvgPool2d(scale),
nn.Conv2d(fc_dim, 512, kernel_size=1, bias=False),
ModuleHelper.BNReLU(512, norm_type=norm_type)))
self.ppm = nn.ModuleList(self.ppm)
def __init__(self,
inplanes,
outplanes,
kernel_size,
stride,
padding=1,
dilation=1,
norm_type=None):
super(_ConvBatchNormReluBlock, self).__init__()
self.conv = nn.Conv2d(in_channels=inplanes,
out_channels=outplanes,
kernel_size=kernel_size,
stride=stride,
padding=padding,
dilation=dilation,
bias=False)
self.bn_relu = ModuleHelper.BNReLU(outplanes, norm_type=norm_type)
def forward(self, data_dict):
out_dict = dict()
label_dict = dict()
loss_dict = dict()
in_img = ModuleHelper.preprocess(data_dict['img'],
self.configer.get('data.normalize'))
out = self.net(in_img)
out_dict['out'] = out
label_dict['out'] = data_dict['label'][:, 0]
if 'ce_loss' in self.valid_loss_dict:
loss_dict['ce_loss'] = dict(params=[out, data_dict['label'][:, 0]],
type=torch.cuda.LongTensor(
[BASE_LOSS_DICT['ce_loss']]),
weight=torch.cuda.FloatTensor(
[self.valid_loss_dict['ce_loss']]))
return out_dict, label_dict, loss_dict
def __init__(self, input_nc, ndf=64, norm_type=None):
"""Construct a 1x1 PatchGAN discriminator
Parameters:
input_nc (int) -- the number of channels in input images
ndf (int) -- the number of filters in the last conv layer
norm_layer -- normalization layer
"""
super(PixelDiscriminator, self).__init__()
use_bias = (norm_type == 'instancenorm')
self.net = [
nn.Conv2d(input_nc, ndf, kernel_size=1, stride=1, padding=0),
nn.LeakyReLU(0.2, True),
nn.Conv2d(ndf, ndf * 2, kernel_size=1, stride=1, padding=0, bias=use_bias),
ModuleHelper.BatchNorm2d(norm_type=norm_type)(ndf * 2),
nn.LeakyReLU(0.2, True),
nn.Conv2d(ndf * 2, 1, kernel_size=1, stride=1, padding=0, bias=use_bias)]
self.net = nn.Sequential(*self.net)
def __init__(self,
features,
out_features=512,
sizes=(1, 2, 3, 6),
norm_type="batchnorm"):
super(PSPModule, self).__init__()
self.stages = []
self.stages = nn.ModuleList([
self._make_stage(features, out_features, size, norm_type)
for size in sizes
])
self.bottleneck = nn.Sequential(
nn.Conv2d(features + len(sizes) * out_features,
out_features,
kernel_size=1,
padding=0,
dilation=1,
bias=False),
ModuleHelper.BNReLU(out_features, norm_type=norm_type),
nn.Dropout2d(0.1))
def __init__(self, inplanes, norm_type="batchnorm", fpn_dim=256):
super(AlignHead, self).__init__()
self.ppm = PSPModule(inplanes,
norm_type=norm_type,
out_features=fpn_dim)
fpn_inplanes = [inplanes // 8, inplanes // 4, inplanes // 2, inplanes]
self.fpn_in = nn.ModuleList()
for fpn_inplane in fpn_inplanes[:-1]:
self.fpn_in.append(
nn.Sequential(
nn.Conv2d(fpn_inplane, fpn_dim, 1),
ModuleHelper.BNReLU(fpn_dim, norm_type=norm_type),
))
self.fpn_out = nn.ModuleList()
self.fpn_out_align = nn.ModuleList()
for i in range(len(fpn_inplanes) - 1):
self.fpn_out.append(
nn.Sequential(conv3x3_bn_relu(fpn_dim, fpn_dim, 1), ))
self.fpn_out_align.append(
AlignModule(inplane=fpn_dim, outplane=fpn_dim // 2))
def __init__(self,
in_channels,
out_channels,
key_channels,
value_channels,
dropout,
sizes=([1]),
norm_type=None,
psp_size=(1, 3, 6, 8)):
super(APNB, self).__init__()
self.stages = []
self.norm_type = norm_type
self.psp_size = psp_size
self.stages = nn.ModuleList([
self._make_stage(in_channels, out_channels, key_channels,
value_channels, size) for size in sizes
])
self.conv_bn_dropout = nn.Sequential(
nn.Conv2d(2 * in_channels, out_channels, kernel_size=1, padding=0),
ModuleHelper.BNReLU(out_channels, norm_type=norm_type),
nn.Dropout2d(dropout))
def __init__(self, configer, loss_dict=None, flag=""):
super(ClsModel, self).__init__()
self.configer = configer
self.flag = flag if len(flag) == 0 else "{}_".format(flag)
self.backbone = ModuleHelper.get_backbone(
backbone=configer.get('network.{}backbone'.format(self.flag)),
pretrained=configer.get('network.{}pretrained'.format(self.flag)))
self.reduction = None
fc_dim_out = configer.get('network.{}fc_dim'.format(self.flag),
default=None)
fc_dim = self.backbone.num_features
if fc_dim_out is not None:
self.reduction = nn.Conv2d(self.backbone.num_features, fc_dim_out,
1)
fc_dim = fc_dim_out
self.linear_list = nn.ModuleList()
linear_type = configer.get('network',
'{}linear_type'.format(self.flag),
default='linear')
self.fc = LINEAR_DICT[linear_type](fc_dim,
configer.get('data.num_classes'))
self.embed = None
if configer.get('network.{}embed'.format(self.flag), default=False):
feat_dim = configer.get('network', '{}feat_dim'.format(self.flag))
self.embed = nn.Sequential(nn.Linear(fc_dim, feat_dim),
nn.BatchNorm1d(feat_dim))
self.bn = None
if configer.get('network.{}bn'.format(self.flag), default=False):
self.bn = nn.BatchNorm1d(fc_dim)
nn.init.zeros_(self.bn.bias)
self.bn.bias.requires_grad = False
self.valid_loss_dict = configer.get(
'loss.loss_weights',
configer.get('loss.loss_type')) if loss_dict is None else loss_dict
def _make_stage(self, features, out_features, size, norm_type):
prior = nn.AdaptiveAvgPool2d(output_size=(size, size))
conv = nn.Conv2d(features, out_features, kernel_size=1, bias=False)
bn = ModuleHelper.BatchNorm2d(norm_type=norm_type)(out_features)
return nn.Sequential(prior, conv, bn)
def __init__(self,
input_nc,
output_nc,
ngf=64,
norm_type=None,
use_dropout=False,
n_blocks=6,
padding_type='reflect'):
"""Construct a Resnet-based generator
Parameters:
input_nc (int) -- the number of channels in input images
output_nc (int) -- the number of channels in output images
ngf (int) -- the number of filters in the last conv layer
norm_layer -- normalization layer
use_dropout (bool) -- if use dropout layers
n_blocks (int) -- the number of ResNet blocks
padding_type (str) -- the name of padding layer in conv layers: reflect | replicate | zero
"""
assert (n_blocks >= 0)
super(ResNetGenerator, self).__init__()
use_bias = (norm_type == 'instancenorm')
model = [
nn.ReflectionPad2d(3),
nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0, bias=use_bias),
ModuleHelper.BatchNorm2d(norm_type=norm_type)(ngf),
nn.ReLU(True)
]
n_downsampling = 2
for i in range(n_downsampling): # add downsampling layers
mult = 2**i
model += [
nn.Conv2d(ngf * mult,
ngf * mult * 2,
kernel_size=3,
stride=2,
padding=1,
bias=use_bias),
ModuleHelper.BatchNorm2d(norm_type=norm_type)(ngf * mult * 2),
nn.ReLU(True)
]
mult = 2**n_downsampling
for i in range(n_blocks): # add ResNet blocks
model += [
ResnetBlock(ngf * mult,
padding_type=padding_type,
norm_type=norm_type,
use_dropout=use_dropout,
use_bias=use_bias)
]
for i in range(n_downsampling): # add upsampling layers
mult = 2**(n_downsampling - i)
model += [
nn.ConvTranspose2d(ngf * mult,
int(ngf * mult / 2),
kernel_size=3,
stride=2,
padding=1,
output_padding=1,
bias=use_bias),
ModuleHelper.BatchNorm2d(norm_type=norm_type)(int(ngf * mult /
2)),
nn.ReLU(True)
]
model += [nn.ReflectionPad2d(3)]
model += [nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0)]
model += [nn.Tanh()]
self.model = nn.Sequential(*model)
def __init__(self,
outer_nc,
inner_nc,
input_nc=None,
submodule=None,
outermost=False,
innermost=False,
norm_type=None,
use_dropout=False):
"""Construct a Unet submodule with skip connections.
Parameters:
outer_nc (int) -- the number of filters in the outer conv layer
inner_nc (int) -- the number of filters in the inner conv layer
input_nc (int) -- the number of channels in input images/features
submodule (UnetSkipConnectionBlock) -- previously defined submodules
outermost (bool) -- if this module is the outermost module
innermost (bool) -- if this module is the innermost module
norm_layer -- normalization layer
user_dropout (bool) -- if use dropout layers.
"""
super(UnetSkipConnectionBlock, self).__init__()
self.outermost = outermost
use_bias = (norm_type == 'instancenorm')
if input_nc is None:
input_nc = outer_nc
downconv = nn.Conv2d(input_nc,
inner_nc,
kernel_size=4,
stride=2,
padding=1,
bias=use_bias)
downrelu = nn.LeakyReLU(0.2, True)
downnorm = ModuleHelper.BatchNorm2d(norm_type=norm_type)(inner_nc)
uprelu = nn.ReLU(True)
upnorm = ModuleHelper.BatchNorm2d(norm_type=norm_type)(outer_nc)
if outermost:
upconv = nn.ConvTranspose2d(inner_nc * 2,
outer_nc,
kernel_size=4,
stride=2,
padding=1)
down = [downconv]
up = [uprelu, upconv, nn.Tanh()]
model = down + [submodule] + up
elif innermost:
upconv = nn.ConvTranspose2d(inner_nc,
outer_nc,
kernel_size=4,
stride=2,
padding=1,
bias=use_bias)
down = [downrelu, downconv]
up = [uprelu, upconv, upnorm]
model = down + up
else:
upconv = nn.ConvTranspose2d(inner_nc * 2,
outer_nc,
kernel_size=4,
stride=2,
padding=1,
bias=use_bias)
down = [downrelu, downconv, downnorm]
up = [uprelu, upconv, upnorm]
if use_dropout:
model = down + [submodule] + up + [nn.Dropout(0.5)]
else:
model = down + [submodule] + up
self.model = nn.Sequential(*model)
