def __init__(self,
input_nc,
output_nc,
nrf,
norm_layer=nn.BatchNorm2d,
use_dropout=False,
n_blocks=6,
padding_type='reflect',
pooling_type='max'):
"""Construct a Resnet-based generator
Parameters:
input_nc (int) -- the number of channels in input images
output_nc (int) -- the number output parameters
nrf (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(Model, self).__init__()
if type(norm_layer) == functools.partial:
use_bias = norm_layer.func == nn.InstanceNorm2d
else:
use_bias = norm_layer == nn.InstanceNorm2d
model = [
nn.ReflectionPad2d(3),
nn.Conv2d(input_nc, nrf, kernel_size=7, padding=0, bias=use_bias),
norm_layer(nrf),
nn.ReLU(True)
]
n_downsampling = 2
for i in range(n_downsampling): # add downsampling layers
mult = 2**i
model += [
nn.Conv2d(nrf * mult,
nrf * mult * 2,
kernel_size=3,
stride=2,
padding=1,
bias=use_bias),
norm_layer(nrf * mult * 2),
nn.ReLU(True)
]
mult = 2**n_downsampling
for i in range(n_blocks): # add ResNet blocks
model += [
ResnetBlock(nrf * mult,
padding_type=padding_type,
norm_layer=norm_layer,
use_dropout=use_dropout,
use_bias=use_bias)
]
# define the regression part
model += [
nn.Conv2d(nrf * mult, nrf * mult, 4, 2,
1), # downsampling spatial, N*64*64*64
norm_layer(nrf * mult),
nn.ReLU(True)
]
if pooling_type == 'max':
model += [nn.AdaptiveMaxPool2d((2, 2))]
elif pooling_type == 'avg':
model += [nn.AdaptiveAvgPool2d((2, 2))]
else:
raise NotImplementedError(f'Unknown pooling type {pooling_type}.')
model += [
# downsampling channels, to N*128*128*128
nn.Conv2d(nrf * mult, nrf * 2, 3, 1, 1),
norm_layer(nrf * 2),
nn.ReLU(True),
# downsampling channels, N*64*128*128
nn.Conv2d(nrf * 2, nrf, 3, 1, 1),
norm_layer(nrf),
nn.ReLU(True),
nn.Flatten(), # flatten the data to shape of N*64
nn.Linear(nrf * 2 * 2, output_nc), # convert to a value
]
self.model = nn.Sequential(*model)
def __init__(self,
input_nc,
output_nc,
ngf=64,
layers=4,
norm='batch',
activation='PReLU',
drop_rate=0,
add_noise=False,
gpu_ids=[],
weight=0.1):
super(_UNetGenerator, self).__init__()
self.gpu_ids = gpu_ids
self.layers = layers
self.weight = weight
norm_layer = get_norm_layer(norm_type=norm)
nonlinearity = get_nonlinearity_layer(activation_type=activation)
if type(norm_layer) == functools.partial:
use_bias = norm_layer.func == nn.InstanceNorm2d
else:
use_bias = norm_layer == nn.InstanceNorm2d
# encoder part
self.pool = nn.AvgPool2d(kernel_size=2, stride=2)
self.conv1 = nn.Sequential(
nn.ReflectionPad2d(3),
nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0, bias=use_bias),
norm_layer(ngf), nonlinearity)
self.conv2 = _EncoderBlock(ngf, ngf * 2, ngf * 2, norm_layer,
nonlinearity, use_bias)
self.conv3 = _EncoderBlock(ngf * 2, ngf * 4, ngf * 4, norm_layer,
nonlinearity, use_bias)
self.conv4 = _EncoderBlock(ngf * 4, ngf * 8, ngf * 8, norm_layer,
nonlinearity, use_bias)
for i in range(layers - 4):
conv = _EncoderBlock(ngf * 8, ngf * 8, ngf * 8, norm_layer,
nonlinearity, use_bias)
setattr(self, 'down' + str(i), conv.model)
center = []
for i in range(7 - layers):
center += [
_InceptionBlock(ngf * 8, ngf * 8, norm_layer, nonlinearity,
7 - layers, drop_rate, use_bias)
]
center += [
_DecoderUpBlock(ngf * 8, ngf * 8, ngf * 4, norm_layer,
nonlinearity, use_bias)
]
if add_noise:
center += [GaussianNoiseLayer()]
self.center = nn.Sequential(*center)
for i in range(layers - 4):
upconv = _DecoderUpBlock(ngf * (8 + 4), ngf * 8, ngf * 4,
norm_layer, nonlinearity, use_bias)
setattr(self, 'up' + str(i), upconv.model)
self.deconv4 = _DecoderUpBlock(ngf * (4 + 4), ngf * 8, ngf * 2,
norm_layer, nonlinearity, use_bias)
self.deconv3 = _DecoderUpBlock(ngf * (2 + 2) + output_nc, ngf * 4, ngf,
norm_layer, nonlinearity, use_bias)
self.deconv2 = _DecoderUpBlock(ngf * (1 + 1) + output_nc, ngf * 2,
int(ngf / 2), norm_layer, nonlinearity,
use_bias)
self.output4 = _OutputBlock(ngf * (4 + 4), output_nc, 3, use_bias)
self.output3 = _OutputBlock(ngf * (2 + 2) + output_nc, output_nc, 3,
use_bias)
self.output2 = _OutputBlock(ngf * (1 + 1) + output_nc, output_nc, 3,
use_bias)
self.output1 = _OutputBlock(
int(ngf / 2) + output_nc, output_nc, 7, use_bias)
self.upsample = nn.Upsample(scale_factor=2, mode='nearest')
def lua_recursive_model(module,seq):
for m in module.modules:
name = type(m).__name__
real = m
if name == 'TorchObject':
name = m._typename.replace('cudnn.','')
m = m._obj
if name == 'SpatialConvolution' or name == 'nn.SpatialConvolutionMM':
if not hasattr(m,'groups') or m.groups is None: m.groups=1
n = nn.Conv2d(m.nInputPlane,m.nOutputPlane,(m.kW,m.kH),(m.dW,m.dH),(m.padW,m.padH),1,m.groups,bias=(m.bias is not None))
copy_param(m,n)
add_submodule(seq,n)
elif name == 'SpatialBatchNormalization':
n = nn.BatchNorm2d(m.running_mean.size(0), m.eps, m.momentum, m.affine)
copy_param(m,n)
add_submodule(seq,n)
elif name == 'VolumetricBatchNormalization':
n = nn.BatchNorm3d(m.running_mean.size(0), m.eps, m.momentum, m.affine)
copy_param(m, n)
add_submodule(seq, n)
elif name == 'ReLU':
n = nn.ReLU()
add_submodule(seq,n)
elif name == 'Sigmoid':
n = nn.Sigmoid()
add_submodule(seq,n)
elif name == 'SpatialMaxPooling':
n = nn.MaxPool2d((m.kW,m.kH),(m.dW,m.dH),(m.padW,m.padH),ceil_mode=m.ceil_mode)
add_submodule(seq,n)
elif name == 'SpatialAveragePooling':
n = nn.AvgPool2d((m.kW,m.kH),(m.dW,m.dH),(m.padW,m.padH),ceil_mode=m.ceil_mode)
add_submodule(seq,n)
elif name == 'SpatialUpSamplingNearest':
n = nn.UpsamplingNearest2d(scale_factor=m.scale_factor)
add_submodule(seq,n)
elif name == 'View':
n = Lambda(lambda x: x.view(x.size(0),-1))
add_submodule(seq,n)
elif name == 'Reshape':
n = Lambda(lambda x: x.view(x.size(0),-1))
add_submodule(seq,n)
elif name == 'Linear':
# Linear in pytorch only accept 2D input
n1 = Lambda(lambda x: x.view(1,-1) if 1==len(x.size()) else x )
n2 = nn.Linear(m.weight.size(1),m.weight.size(0),bias=(m.bias is not None))
copy_param(m,n2)
n = nn.Sequential(n1,n2)
add_submodule(seq,n)
elif name == 'Dropout':
m.inplace = False
n = nn.Dropout(m.p)
add_submodule(seq,n)
elif name == 'SoftMax':
n = nn.Softmax()
add_submodule(seq,n)
elif name == 'Identity':
n = Lambda(lambda x: x) # do nothing
add_submodule(seq,n)
elif name == 'SpatialFullConvolution':
n = nn.ConvTranspose2d(m.nInputPlane,m.nOutputPlane,(m.kW,m.kH),(m.dW,m.dH),(m.padW,m.padH),(m.adjW,m.adjH))
copy_param(m,n)
add_submodule(seq,n)
elif name == 'VolumetricFullConvolution':
n = nn.ConvTranspose3d(m.nInputPlane,m.nOutputPlane,(m.kT,m.kW,m.kH),(m.dT,m.dW,m.dH),(m.padT,m.padW,m.padH),(m.adjT,m.adjW,m.adjH),m.groups)
copy_param(m,n)
add_submodule(seq, n)
elif name == 'SpatialReplicationPadding':
n = nn.ReplicationPad2d((m.pad_l,m.pad_r,m.pad_t,m.pad_b))
add_submodule(seq,n)
elif name == 'SpatialReflectionPadding':
n = nn.ReflectionPad2d((m.pad_l,m.pad_r,m.pad_t,m.pad_b))
add_submodule(seq,n)
elif name == 'Copy':
n = Lambda(lambda x: x) # do nothing
add_submodule(seq,n)
elif name == 'Narrow':
n = Lambda(lambda x,a=(m.dimension,m.index,m.length): x.narrow(*a))
add_submodule(seq,n)
elif name == 'SpatialCrossMapLRN':
lrn = lnn.SpatialCrossMapLRN(m.size,m.alpha,m.beta,m.k)
n = Lambda(lambda x,lrn=lrn: Variable(lrn.forward(x.data)))
add_submodule(seq,n)
elif name == 'Sequential':
n = nn.Sequential()
lua_recursive_model(m,n)
add_submodule(seq,n)
elif name == 'ConcatTable': # output is list
n = LambdaMap(lambda x: x)
lua_recursive_model(m,n)
add_submodule(seq,n)
elif name == 'CAddTable': # input is list
n = LambdaReduce(lambda x,y: x+y)
add_submodule(seq,n)
elif name == 'Concat':
dim = m.dimension
n = LambdaReduce(lambda x,y,dim=dim: torch.cat((x,y),dim))
lua_recursive_model(m,n)
add_submodule(seq,n)
elif name == 'TorchObject':
print('Not Implement',name,real._typename)
else:
print('Not Implement',name)
def __init__(self, ngf=64, img_size=256, light=False):
super(ResnetGenerator, self).__init__()
self.light = light
self.ConvBlock1 = nn.Sequential(
nn.ReflectionPad2d(3),
nn.Conv2d(3, ngf, kernel_size=7, stride=1, padding=0, bias=False),
nn.InstanceNorm2d(ngf), nn.ReLU(True))
self.HourGlass1 = HourGlass(ngf, ngf)
self.HourGlass2 = HourGlass(ngf, ngf)
# Down-Sampling
self.DownBlock1 = nn.Sequential(
nn.ReflectionPad2d(1),
nn.Conv2d(ngf,
ngf * 2,
kernel_size=3,
stride=2,
padding=0,
bias=False), nn.InstanceNorm2d(ngf * 2), nn.ReLU(True))
self.DownBlock2 = nn.Sequential(
nn.ReflectionPad2d(1),
nn.Conv2d(ngf * 2,
ngf * 4,
kernel_size=3,
stride=2,
padding=0,
bias=False), nn.InstanceNorm2d(ngf * 4), nn.ReLU(True))
# Encoder Bottleneck
self.EncodeBlock1 = ResnetBlock(ngf * 4)
self.EncodeBlock2 = ResnetBlock(ngf * 4)
self.EncodeBlock3 = ResnetBlock(ngf * 4)
self.EncodeBlock4 = ResnetBlock(ngf * 4)
# Class Activation Map
self.gap_fc = nn.Linear(ngf * 4, 1)
self.gmp_fc = nn.Linear(ngf * 4, 1)
self.conv1x1 = nn.Conv2d(ngf * 8, ngf * 4, kernel_size=1, stride=1)
self.relu = nn.ReLU(True)
# Gamma, Beta block
if self.light:
self.FC = nn.Sequential(nn.Linear(ngf * 4, ngf * 4), nn.ReLU(True),
nn.Linear(ngf * 4, ngf * 4), nn.ReLU(True))
else:
self.FC = nn.Sequential(
nn.Linear(img_size // 4 * img_size // 4 * ngf * 4, ngf * 4),
nn.ReLU(True), nn.Linear(ngf * 4, ngf * 4), nn.ReLU(True))
# Decoder Bottleneck
self.DecodeBlock1 = ResnetSoftAdaLINBlock(ngf * 4)
self.DecodeBlock2 = ResnetSoftAdaLINBlock(ngf * 4)
self.DecodeBlock3 = ResnetSoftAdaLINBlock(ngf * 4)
self.DecodeBlock4 = ResnetSoftAdaLINBlock(ngf * 4)
# Up-Sampling
self.UpBlock1 = nn.Sequential(
nn.Upsample(scale_factor=2), nn.ReflectionPad2d(1),
nn.Conv2d(ngf * 4,
ngf * 2,
kernel_size=3,
stride=1,
padding=0,
bias=False), LIN(ngf * 2), nn.ReLU(True))
self.UpBlock2 = nn.Sequential(
nn.Upsample(scale_factor=2), nn.ReflectionPad2d(1),
nn.Conv2d(ngf * 2,
ngf,
kernel_size=3,
stride=1,
padding=0,
bias=False), LIN(ngf), nn.ReLU(True))
self.HourGlass3 = HourGlass(ngf, ngf)
self.HourGlass4 = HourGlass(ngf, ngf, False)
self.ConvBlock2 = nn.Sequential(
nn.ReflectionPad2d(3),
nn.Conv2d(3, 3, kernel_size=7, stride=1, padding=0, bias=False),
nn.Tanh())
def __init__(self,
input_nc,
output_nc,
ngf=64,
norm_layer=nn.BatchNorm2d,
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__()
if type(norm_layer) == functools.partial:
use_bias = norm_layer.func == nn.InstanceNorm2d
else:
use_bias = norm_layer == nn.InstanceNorm2d
model = [
nn.ReflectionPad2d(3),
nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0, bias=use_bias),
norm_layer(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),
norm_layer(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_layer=norm_layer,
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),
norm_layer(int(ngf * mult / 2)),
nn.ReLU(True)
]
# model += [nn.Upsample(scale_factor = 2, mode='bilinear'),
# nn.ReflectionPad2d(1),
# nn.Conv2d(ngf * mult, int(ngf * mult / 2),
# kernel_size=3, stride=1, padding=0),
# norm_layer(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,
down_sample_nl,
input_nc,
output_nc,
ngf=64,
norm_layer=nn.BatchNorm2d,
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, first???
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__()
if type(norm_layer) == functools.partial:
use_bias = norm_layer.func == nn.InstanceNorm2d
else:
use_bias = norm_layer == nn.InstanceNorm2d
model = OrderedDict()
model['rgb2channel_pad'] = nn.ReflectionPad2d(3)
model['rgb2channel_conv'] = nn.Conv2d(input_nc,
ngf,
kernel_size=7,
padding=0,
bias=use_bias)
model['rgb2channel_norm'] = norm_layer(ngf)
model['rgb2channel_relu'] = nn.ReLU(True)
#model = [nn.ReflectionPad2d(3),
# nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0, bias=use_bias),
# norm_layer(ngf),
# nn.ReLU(True)]
n_downsampling = down_sample_nl
for i in range(n_downsampling): # add downsampling layers
mult = 2**i
#original
model[str('down_sample_conv_' +
str(n_downsampling - i - 1))] = nn.Conv2d(ngf * mult,
ngf * mult * 2,
kernel_size=3,
stride=2,
padding=1,
bias=use_bias)
model[str('down_sample_norm_' +
str(n_downsampling - i - 1))] = norm_layer(ngf * mult *
2)
model[str('down_sample_relu_' +
str(n_downsampling - i - 1))] = nn.ReLU(True)
#model += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1, bias=use_bias),
# norm_layer(ngf * mult * 2),
# nn.ReLU(True)]
mult = 2**n_downsampling
for i in range(n_blocks): # add ResNet blocks
model[str('resnet_' + str(i))] = ResnetBlock(
ngf * mult,
padding_type=padding_type,
norm_layer=norm_layer,
use_dropout=use_dropout,
use_bias=use_bias)
#model += [ResnetBlock(ngf * mult, padding_type=padding_type, norm_layer=norm_layer, use_dropout=use_dropout, use_bias=use_bias)]
for i in range(n_downsampling): # add upsampling layers
mult = 2**(n_downsampling - i)
model[str('up_sample_conv_' + str(i))] = nn.ConvTranspose2d(
ngf * mult,
int(ngf * mult / 2),
kernel_size=3,
stride=2,
padding=1,
output_padding=1,
bias=use_bias)
model[str('up_sample_norm_' + str(i))] = norm_layer(
int(ngf * mult / 2))
model[str('up_sample_relu_' + str(i))] = nn.ReLU(True)
#model += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2),
# kernel_size=3, stride=2,
# padding=1, output_padding=1,
# bias=use_bias),
# norm_layer(int(ngf * mult / 2)),
# nn.ReLU(True)]
model['channel2rgb_pad'] = nn.ReflectionPad2d(3)
model['channel2rgb_conv'] = nn.Conv2d(ngf,
output_nc,
kernel_size=7,
padding=0)
model['channel2rgb_tanh'] = nn.Tanh()
#model['channle2rgb_pad'] += [nn.ReflectionPad2d(3)]
#model += [nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0)]
#model += [nn.Tanh()]
#self.model = nn.Sequential(*model)
self.model = nn.Sequential(model)
def __init__(self):
super(level_5_decoder, self).__init__()
d = load_t7("models/feature_invertor_conv5_1.t7")
# decoder
self.reflecPad15 = nn.ReflectionPad2d((1, 1, 1, 1))
self.conv15 = nn.Conv2d(512, 512, 3, 1, 0)
self.conv15.weight = torch.nn.Parameter(
torch.from_numpy(d[0]["w"]).cuda())
self.conv15.bias = torch.nn.Parameter(
torch.from_numpy(d[0]["b"]).cuda())
self.relu15 = nn.ReLU(inplace=True)
self.unpool = nn.UpsamplingNearest2d(scale_factor=2)
# 28 x 28
self.reflecPad16 = nn.ReflectionPad2d((1, 1, 1, 1))
self.conv16 = nn.Conv2d(512, 512, 3, 1, 0)
self.conv16.weight = torch.nn.Parameter(
torch.from_numpy(d[1]["w"]).cuda())
self.conv16.bias = torch.nn.Parameter(
torch.from_numpy(d[1]["b"]).cuda())
self.relu16 = nn.ReLU(inplace=True)
# 28 x 28
self.reflecPad17 = nn.ReflectionPad2d((1, 1, 1, 1))
self.conv17 = nn.Conv2d(512, 512, 3, 1, 0)
self.conv17.weight = torch.nn.Parameter(
torch.from_numpy(d[2]["w"]).cuda())
self.conv17.bias = torch.nn.Parameter(
torch.from_numpy(d[2]["b"]).cuda())
self.relu17 = nn.ReLU(inplace=True)
# 28 x 28
self.reflecPad18 = nn.ReflectionPad2d((1, 1, 1, 1))
self.conv18 = nn.Conv2d(512, 512, 3, 1, 0)
self.conv18.weight = torch.nn.Parameter(
torch.from_numpy(d[3]["w"]).cuda())
self.conv18.bias = torch.nn.Parameter(
torch.from_numpy(d[3]["b"]).cuda())
self.relu18 = nn.ReLU(inplace=True)
# 28 x 28
self.reflecPad19 = nn.ReflectionPad2d((1, 1, 1, 1))
self.conv19 = nn.Conv2d(512, 256, 3, 1, 0)
self.conv19.weight = torch.nn.Parameter(
torch.from_numpy(d[4]["w"]).cuda())
self.conv19.bias = torch.nn.Parameter(
torch.from_numpy(d[4]["b"]).cuda())
self.relu19 = nn.ReLU(inplace=True)
# 28 x 28
self.unpool2 = nn.UpsamplingNearest2d(scale_factor=2)
# 56 x 56
self.reflecPad20 = nn.ReflectionPad2d((1, 1, 1, 1))
self.conv20 = nn.Conv2d(256, 256, 3, 1, 0)
self.conv20.weight = torch.nn.Parameter(
torch.from_numpy(d[5]["w"]).cuda())
self.conv20.bias = torch.nn.Parameter(
torch.from_numpy(d[5]["b"]).cuda())
self.relu20 = nn.ReLU(inplace=True)
# 56 x 56
self.reflecPad21 = nn.ReflectionPad2d((1, 1, 1, 1))
self.conv21 = nn.Conv2d(256, 256, 3, 1, 0)
self.conv21.weight = torch.nn.Parameter(
torch.from_numpy(d[6]["w"]).cuda())
self.conv21.bias = torch.nn.Parameter(
torch.from_numpy(d[6]["b"]).cuda())
self.relu21 = nn.ReLU(inplace=True)
self.reflecPad22 = nn.ReflectionPad2d((1, 1, 1, 1))
self.conv22 = nn.Conv2d(256, 256, 3, 1, 0)
self.conv22.weight = torch.nn.Parameter(
torch.from_numpy(d[7]["w"]).cuda())
self.conv22.bias = torch.nn.Parameter(
torch.from_numpy(d[7]["b"]).cuda())
self.relu22 = nn.ReLU(inplace=True)
self.reflecPad23 = nn.ReflectionPad2d((1, 1, 1, 1))
self.conv23 = nn.Conv2d(256, 128, 3, 1, 0)
self.conv23.weight = torch.nn.Parameter(
torch.from_numpy(d[8]["w"]).cuda())
self.conv23.bias = torch.nn.Parameter(
torch.from_numpy(d[8]["b"]).cuda())
self.relu23 = nn.ReLU(inplace=True)
self.unpool3 = nn.UpsamplingNearest2d(scale_factor=2)
# 112 X 112
self.reflecPad24 = nn.ReflectionPad2d((1, 1, 1, 1))
self.conv24 = nn.Conv2d(128, 128, 3, 1, 0)
self.conv24.weight = torch.nn.Parameter(
torch.from_numpy(d[9]["w"]).cuda())
self.conv24.bias = torch.nn.Parameter(
torch.from_numpy(d[9]["b"]).cuda())
self.relu24 = nn.ReLU(inplace=True)
self.reflecPad25 = nn.ReflectionPad2d((1, 1, 1, 1))
self.conv25 = nn.Conv2d(128, 64, 3, 1, 0)
self.conv25.weight = torch.nn.Parameter(
torch.from_numpy(d[10]["w"]).cuda())
self.conv25.bias = torch.nn.Parameter(
torch.from_numpy(d[10]["b"]).cuda())
self.relu25 = nn.ReLU(inplace=True)
self.unpool4 = nn.UpsamplingNearest2d(scale_factor=2)
self.reflecPad26 = nn.ReflectionPad2d((1, 1, 1, 1))
self.conv26 = nn.Conv2d(64, 64, 3, 1, 0)
self.conv26.weight = torch.nn.Parameter(
torch.from_numpy(d[11]["w"]).cuda())
self.conv26.bias = torch.nn.Parameter(
torch.from_numpy(d[11]["b"]).cuda())
self.relu26 = nn.ReLU(inplace=True)
self.reflecPad27 = nn.ReflectionPad2d((1, 1, 1, 1))
self.conv27 = nn.Conv2d(64, 3, 3, 1, 0)
self.conv27.weight = torch.nn.Parameter(
torch.from_numpy(d[12]["w"]).cuda())
self.conv27.bias = torch.nn.Parameter(
torch.from_numpy(d[12]["b"]).cuda())
def __init__(self, input_nc, output_nc, n_residual_blocks=9):
"""
定义生成网络
参数：
input_nc --输入通道数
output_nc --输出通道数
n_residual_blocks --残差模块数量
"""
super(Generator, self).__init__()
# 初始化卷积模块
# 因为使用ReflectionPad扩充
# 所以输入是3*256*256
# 输出是64*256*256
model = [
nn.ReflectionPad2d(3),
nn.Conv2d(input_nc, 64, 7),
nn.InstanceNorm2d(64),
nn.ReLU(inplace=True)
]
# 进行下采样
# 第一个range：输入是64*256*256，输出是128*128*128
# 第二个range：输入是128*128*128，输出是256*64*64
in_features = 64
out_features = in_features * 2
for _ in range(2):
model += [
nn.Conv2d(in_features, out_features, 3, stride=2, padding=1),
nn.InstanceNorm2d(out_features),
nn.ReLU(inplace=True)
]
in_features = out_features
out_features = in_features * 2
# 使用残差模块
# 输入输出都是256*64*64
for _ in range(n_residual_blocks): # 默认添加9个残差模块
model += [ResidualBlock(in_features)]
# 进行上采样
# 第一个range：输入是256*64*64，输出是128*128*128
# 第二个range：输入是128*128*128，输出是64*256*256
out_features = in_features // 2
for _ in range(2):
model += [
nn.ConvTranspose2d(in_features,
out_features,
3,
stride=2,
padding=1,
output_padding=1),
nn.InstanceNorm2d(out_features),
nn.ReLU(inplace=True)
]
in_features = out_features
out_features = in_features // 2
# 最后输出层
# 输入是64*256*256
# 输出是3*256*256
model += [
nn.ReflectionPad2d(3),
nn.Conv2d(64, output_nc, 7),
nn.Tanh()
]
self.model = nn.Sequential(*model)
#>一 镜像 padding ##>1.1 一维镜像padding m1 = nn.ReflectionPad1d(2) ##left=right=2 m2 = nn.ReflectionPad1d((2, 1)) ##left=2,right=1 input = torch.arange(8, dtype=torch.float).reshape(1, 2, 4) input #%% m1(input) #%% m2(input) #%% ##>1.2 二维镜像padding m1 = nn.ReflectionPad2d(2) #left=right=top=bottom=2 m2 = nn.ReflectionPad2d((1, 1, 2, 0)) #left=right=1,top=2,bottom=0 input = torch.arange(9, dtype=torch.float).reshape(1, 1, 3, 3) input #%% m1(input) #%% m2(input) #%% #>二 复制 padding ##>2.1 一维复制padding m1 = nn.ReplicationPad1d(2) m2 = nn.ReplicationPad1d((2, 1))
relu3_1 = self.relu3_1(intr2)
if layer == 'relu3_1':
return crelu3_1
intr3 = F.max_pool2d(self.intr3(relu3_1),
kernel_size=2,
stride=2,
padding=0)
relu4_1 = self.relu4_1(intr3)
if layer == 'relu4_1':
return relu4_1
vgg_output = namedtuple('vgg_output',
['relu1_1', 'relu2_1', 'relu3_1', 'relu4_1'])
return vgg_output(relu1_1, relu2_1, relu3_1, relu4_1)
vgg_decoder_relu5_1 = nn.Sequential(nn.ReflectionPad2d((1, 1, 1, 1)),
nn.Conv2d(512, 512, 3), nn.ReLU(),
nn.Upsample(scale_factor=2),
nn.ReflectionPad2d((1, 1, 1, 1)),
nn.Conv2d(512, 512, 3), nn.ReLU(),
nn.ReflectionPad2d((1, 1, 1, 1)),
nn.Conv2d(512, 512, 3), nn.ReLU(),
nn.ReflectionPad2d((1, 1, 1, 1)),
nn.Conv2d(512, 512, 3), nn.ReLU(),
nn.ReflectionPad2d((1, 1, 1, 1)),
nn.Conv2d(512, 256, 3), nn.ReLU(),
nn.Upsample(scale_factor=2),
nn.ReflectionPad2d((1, 1, 1, 1)),
nn.Conv2d(256, 256, 3), nn.ReLU(),
nn.ReflectionPad2d((1, 1, 1, 1)),
nn.Conv2d(256, 256, 3), nn.ReLU(),
def __init__(self, opt):
super().__init__()
input_nc = opt.label_nc + (1 if opt.contain_dontcare_label else
0) + (1 if opt.no_instance else 1)
norm_layer = get_nonspade_norm_layer(opt, opt.norm_G)
activation = nn.ReLU(False)
model = []
# initial conv
model += [
nn.ReflectionPad2d(opt.resnet_initial_kernel_size // 2),
norm_layer(
nn.Conv2d(input_nc,
opt.ngf,
kernel_size=opt.resnet_initial_kernel_size,
padding=0)), activation
]
# downsample
mult = 1
for i in range(opt.resnet_n_downsample):
model += [
norm_layer(
nn.Conv2d(opt.ngf * mult,
opt.ngf * mult * 2,
kernel_size=3,
stride=2,
padding=1)), activation
]
mult *= 2
# resnet blocks
for i in range(opt.resnet_n_blocks):
model += [
ResnetBlock(opt.ngf * mult,
norm_layer=norm_layer,
activation=activation,
kernel_size=opt.resnet_kernel_size)
]
# upsample
for i in range(opt.resnet_n_downsample):
nc_in = int(opt.ngf * mult)
nc_out = int((opt.ngf * mult) / 2)
model += [
norm_layer(
nn.ConvTranspose2d(nc_in,
nc_out,
kernel_size=3,
stride=2,
padding=1,
output_padding=1)), activation
]
mult = mult // 2
# final output conv
model += [
nn.ReflectionPad2d(3),
nn.Conv2d(nc_out, opt.output_nc, kernel_size=7, padding=0),
nn.Tanh()
]
self.model = nn.Sequential(*model)
def train(model, train_loader, test_loader, mode='EDSR_Baseline', save_image_every=50, save_model_every=10, test_model_every=1, epoch_start=0, num_epochs=1000, device=None, refresh=True, scale=2, today=None):
if device is None:
device = 'cuda' if torch.cuda.is_available() else 'cpu'
if today is None: today = datetime.datetime.now().strftime('%Y.%m.%d')
result_dir = f'./results/{today}/{mode}'
weight_dir = f'./weights/{today}/{mode}'
logger_dir = f'./logger/{today}_{mode}'
csv = f'./hist_{today}_{mode}.csv'
if refresh:
try:
shutil.rmtree(result_dir)
shutil.rmtree(weight_dir)
shutil.rmtree(logger_dir)
except FileNotFoundError:
pass
os.makedirs(result_dir, exist_ok=True)
os.makedirs(weight_dir, exist_ok=True)
os.makedirs(logger_dir, exist_ok=True)
logger = SummaryWriter(log_dir=logger_dir, flush_secs=2)
model = model.to(device)
# model.load_state_dict(torch.load('./weights/2021.03.10/HMNET_x4_Heavy_no_fea_REDS_size_0/epoch_0022.pth'))
params = list(model.parameters())
optim = torch.optim.Adam(params, lr=1e-4)
scheduler = torch.optim.lr_scheduler.StepLR(optim, step_size=1000, gamma= 0.99)
criterion = torch.nn.MSELoss()
GMSD = GMSD_quality().to(device)
opening = Opening().to(device)
blur = Blur().to(device)
mshf = MSHF(3, 3).to(device)
high_pass = high_pass_filter(scale=4).to(device).eval()
high_pass.load_state_dict(torch.load('./weights/2021.02.01/HMNET_x4_High_Pass_Filter_lite/epoch_0100.pth'))
downx2_bicubic = nn.Upsample(scale_factor=1/2, mode='bicubic', align_corners=False)
downx4_bicubic = nn.Upsample(scale_factor=1/4, mode='bicubic', align_corners=False)
padl = nn.ReflectionPad2d(10)
padh = nn.ReflectionPad2d(40)
start_time = time.time()
print(f'Training Start || Mode: {mode}')
step = 0
pfix = OrderedDict()
pfix_test = OrderedDict()
hist = dict()
hist['mode'] = f'{today}_{mode}'
for key in ['epoch', 'psnr', 'ssim', 'ms-ssim']:
hist[key] = []
soft_mask = False
# hf_kernel = get_hf_kernel(mode='high')
for epoch in range(epoch_start, epoch_start+num_epochs):
if epoch == 0:
torch.save(model.state_dict(), f'{weight_dir}/epoch_{epoch+1:04d}.pth')
if epoch == 0:
with torch.no_grad():
with tqdm(test_loader, desc=f'{mode} || Warming Up || Test Epoch {epoch}/{num_epochs}', position=0, leave=True) as pbar_test:
psnrs = []
ssims = []
msssims = []
for lr, hr, fname in pbar_test:
lr = lr.to(device)
hr = hr.to(device)
sr, srx2, srx1 = model(lr)
sr = quantize(sr)
psnr, ssim, msssim = evaluate(hr, sr)
psnrs.append(psnr)
ssims.append(ssim)
msssims.append(msssim)
psnr_mean = np.array(psnrs).mean()
ssim_mean = np.array(ssims).mean()
msssim_mean = np.array(msssims).mean()
pfix_test['psnr_mean'] = f'{psnr_mean:.4f}'
pfix_test['ssim_mean'] = f'{ssim_mean:.4f}'
pfix_test['msssim_mean'] = f'{msssim_mean:.4f}'
pbar_test.set_postfix(pfix_test)
if len(psnrs) > 1: break
with tqdm(train_loader, desc=f'{mode} || Epoch {epoch+1}/{num_epochs}', position=0, leave=True) as pbar:
psnrs = []
ssims = []
msssims = []
losses = []
for lr, hr, _ in pbar:
lr = lr.to(device)
hr = hr.to(device)
lr = padl(lr)
hr = padh(hr)
#hrx1 = downx4_bicubic(hr)
#hrx2 = downx2_bicubic(hr)
# prediction
sr, srx2, srx1 = model(lr)
#sr = sr[:,:,40:-40,40:-40]
gmsd = GMSD(hr, sr)
sr_ = quantize(sr)
psnr, ssim, msssim = evaluate(hr, sr_)
#if psnr >= 40 - 2*scale:
# soft_mask = True
#else:
soft_mask = False
if soft_mask:
# with torch.no_grad():
# for _ in range(10): gmsd = opening(gmsd)
gmask = gmsd / gmsd.max()
gmask = (gmask > 0.2) * 1.0
gmask = blur(gmask)
gmask = (gmask - gmask.min()) / (gmask.max() - gmask.min() + 1e-7)
gmask = (gmask + 0.25) / 1.25
gmask = gmask.detach()
gmaskx2 = downx2_bicubic(gmask)
gmaskx1 = downx4_bicubic(gmask)
# training
loss = criterion(sr * gmask, hr * gmask)
lossx2 = criterion(srx2 * gmaskx2, hrx2 * gmaskx2)
lossx1 = criterion(srx1 * gmaskx1, hrx1 * gmaskx1)
else:
loss = criterion(sr, hr)
#lossx2 = criterion(srx2, hrx2)
#lossx1 = criterion(srx1, hrx1)
loss_hf = criterion(high_pass(sr), high_pass(hr))
##loss_hf += 0.25 * criterion(high_pass(srx2), high_pass(hrx2))
#loss_hf += 0.125 * criterion(high_pass(srx1), high_pass(hrx1))
# training
loss_tot = loss + loss_hf #+ 0.25 * lossx2 + 0.125 * lossx1 + loss_hf * 0.2
optim.zero_grad()
loss_tot.backward()
optim.step()
scheduler.step()
# training history
elapsed_time = time.time() - start_time
elapsed = sec2time(elapsed_time)
pfix['Step'] = f'{step+1}'
pfix['Loss'] = f'{loss_tot.item():.4f}'
# pfix['Loss x2'] = f'{lossx2.item():.4f}'
# pfix['Loss x1'] = f'{lossx1.item():.4f}'
pfix['Elapsed'] = f'{elapsed}'
psnrs.append(psnr)
ssims.append(ssim)
msssims.append(msssim)
psnr_mean = np.array(psnrs).mean()
ssim_mean = np.array(ssims).mean()
msssim_mean = np.array(msssims).mean()
pfix['PSNR_mean'] = f'{psnr_mean:.2f}'
pfix['SSIM_mean'] = f'{ssim_mean:.4f}'
free_gpu = get_gpu_memory()[0]
pfix['free GPU'] = f'{free_gpu}MiB'
pbar.set_postfix(pfix)
losses.append(loss.item())
if step % save_image_every == 0:
z = torch.zeros_like(lr[0])
_, _, llr, _ = lr.shape
_, _, hlr, _ = hr.shape
if hlr // 2 == llr:
xz = torch.cat((lr[0], z), dim=-2)
elif hlr // 4 == llr:
xz = torch.cat((lr[0], z, z, z), dim=-2)
imsave([xz, sr[0], hr[0], gmsd[0]], f'{result_dir}/epoch_{epoch+1}_iter_{step:05d}.jpg')
step += 1
logger.add_scalar("Loss/train", np.array(losses).mean(), epoch+1)
logger.add_scalar("PSNR/train", psnr_mean, epoch+1)
logger.add_scalar("SSIM/train", ssim_mean, epoch+1)
if (epoch+1) % save_model_every == 0:
torch.save(model.state_dict(), f'{weight_dir}/epoch_{epoch+1:04d}.pth')
if (epoch+1) % test_model_every == 0:
with torch.no_grad():
with tqdm(test_loader, desc=f'{mode} || Test Epoch {epoch+1}/{num_epochs}', position=0, leave=True) as pbar_test:
psnrs = []
ssims = []
msssims = []
for lr, hr, fname in pbar_test:
fname = fname[0].split('/')[-1].split('.pt')[0]
lr = lr.to(device)
hr = hr.to(device)
sr, _, _ = model(lr)
mshf_hr = mshf(hr)
mshf_sr = mshf(sr)
gmsd = GMSD(hr, sr)
sr = quantize(sr)
psnr, ssim, msssim = evaluate(hr, sr)
psnrs.append(psnr)
ssims.append(ssim)
msssims.append(msssim)
psnr_mean = np.array(psnrs).mean()
ssim_mean = np.array(ssims).mean()
msssim_mean = np.array(msssims).mean()
pfix_test['psnr_mean'] = f'{psnr_mean:.4f}'
pfix_test['ssim_mean'] = f'{ssim_mean:.4f}'
pfix_test['msssim_mean'] = f'{msssim_mean:.4f}'
pbar_test.set_postfix(pfix_test)
z = torch.zeros_like(lr[0])
_, _, llr, _ = lr.shape
_, _, hlr, _ = hr.shape
if hlr // 2 == llr:
xz = torch.cat((lr[0], z), dim=-2)
elif hlr // 4 == llr:
xz = torch.cat((lr[0], z, z, z), dim=-2)
imsave([xz, sr[0], hr[0], gmsd[0]], f'{result_dir}/{fname}.jpg')
mshf_vis = torch.cat((torch.cat([mshf_sr[:,i,:,:] for i in range(mshf_sr.shape[1])], dim=-1),
torch.cat([mshf_hr[:,i,:,:] for i in range(mshf_hr.shape[1])], dim=-1)), dim=-2)
imsave(mshf_vis, f'{result_dir}/MSHF_{fname}.jpg')
hist['epoch'].append(epoch+1)
hist['psnr'].append(psnr_mean)
hist['ssim'].append(ssim_mean)
hist['ms-ssim'].append(msssim_mean)
logger.add_scalar("PSNR/test", psnr_mean, epoch+1)
logger.add_scalar("SSIM/test", ssim_mean, epoch+1)
logger.add_scalar("MS-SSIM/test", msssim_mean, epoch+1)
df = pd.DataFrame(hist)
df.to_csv(csv)
def __init__(self,
outer_nc,
inner_nc,
input_nc=None,
submodule=None,
outermost=False,
innermost=False,
norm_layer=nn.BatchNorm2d,
use_dropout=False,
use_transpose=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
use_dropout (bool) -- if use dropout layers.
"""
super(UnetSkipConnectionBlock, self).__init__()
self.outermost = outermost
if type(norm_layer) == functools.partial:
use_bias = norm_layer.func == nn.InstanceNorm2d
else:
use_bias = norm_layer == nn.InstanceNorm2d
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 = norm_layer(inner_nc)
uprelu = nn.ReLU(True)
upnorm = norm_layer(outer_nc)
if outermost:
upsample = nn.Upsample(scale_factor=2, mode='bilinear')
reflect = nn.ReflectionPad2d(1)
upconv = nn.Conv2d(inner_nc * 2,
outer_nc,
kernel_size=3,
stride=1,
padding=0)
#upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc,
# kernel_size=4, stride=2,
# padding=1)
down = [downconv]
up = [uprelu, upsample, reflect, 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)
if use_transpose:
upconv = nn.ConvTranspose2d(inner_nc * 2,
outer_nc,
kernel_size=4,
stride=2,
padding=1)
up = [uprelu, upconv, upnorm]
else:
upsample = nn.Upsample(scale_factor=2, mode='nearest')
reflect = nn.ReflectionPad2d(1)
upconv = nn.Conv2d(inner_nc * 2,
outer_nc,
kernel_size=3,
stride=1,
padding=0)
up = [uprelu, upsample, reflect, upconv, upnorm]
down = [downrelu, downconv, downnorm]
# down = [downrelu, downconv, downnorm]
if use_dropout:
model = down + [submodule] + up + [nn.Dropout(0.5)]
else:
model = down + [submodule] + up
self.model = nn.Sequential(*model)
def __init__(self):
super(STNBasic, self).__init__()
self.features1 = nn.Sequential(
nn.ReflectionPad2d((1, 1, 1, 1)),
nn.Conv2d(3, 16, kernel_size=3),
nn.BatchNorm2d(16),
nn.ReLU(),
nn.ReflectionPad2d((1, 1, 1, 1)),
nn.Conv2d(16, 16, kernel_size=3),
nn.BatchNorm2d(16),
nn.ReLU(),
nn.MaxPool2d(2), ######
nn.ReflectionPad2d((1, 1, 1, 1)),
nn.Conv2d(16, 32, kernel_size=3),
nn.BatchNorm2d(32),
nn.ReLU(),
nn.ReflectionPad2d((1, 1, 1, 1)),
nn.Conv2d(32, 32, kernel_size=3),
nn.BatchNorm2d(32),
nn.ReLU(),
nn.MaxPool2d(2), ######
)
self.features2 = nn.Sequential(
nn.ReflectionPad2d((1, 1, 1, 1)),
nn.Conv2d(32, 32, kernel_size=3),
nn.BatchNorm2d(32),
nn.ReLU(),
nn.ReflectionPad2d((1, 1, 1, 1)),
nn.Conv2d(32, 32, kernel_size=3),
nn.BatchNorm2d(32),
nn.ReLU(),
nn.MaxPool2d(2), ######
)
self.regressor1 = nn.Sequential(nn.BatchNorm1d(8 * 8 * 32),
nn.Dropout(p=0.6), nn.ReLU(),
nn.Linear(8 * 8 * 32, 30),
nn.Dropout(p=0.6), nn.Linear(30, 2))
self.regressor2 = nn.Sequential(nn.BatchNorm1d(6), nn.Dropout(p=0.4),
nn.ReLU(), nn.Linear(6, 10),
nn.Dropout(p=0.4), nn.Linear(10, 2))
self.regressor3 = nn.Sequential(nn.Linear(4, 10), nn.Linear(10, 2))
# Spatial transformer localization-network
self.localization = nn.Sequential(
nn.ReflectionPad2d((1, 1, 1, 1)),
nn.Conv2d(3, 16, kernel_size=3),
nn.BatchNorm2d(16),
nn.ReLU(),
nn.MaxPool2d(2), ####
nn.ReflectionPad2d((1, 1, 1, 1)),
nn.Conv2d(16, 32, kernel_size=3),
nn.BatchNorm2d(32),
nn.ReLU(),
nn.MaxPool2d(2), ####
nn.ReflectionPad2d((1, 1, 1, 1)),
nn.Conv2d(32, 32, kernel_size=3),
nn.BatchNorm2d(32),
nn.ReLU(),
nn.MaxPool2d(2) ####
)
# Regressor for the 3 * 2 affine matrix
self.fc_loc = nn.Sequential(nn.Linear(32 * 8 * 8, 20), nn.ReLU(),
nn.Linear(20, 3 * 2))
# Initialize the weights/bias with identity transformation
self.fc_loc[2].weight.data.zero_()
self.fc_loc[2].bias.data.copy_(
torch.tensor([1, 0, 0, 0, 1, 0], dtype=torch.float))
def test_pad(self):
x = torch.tensor([[[[0.0, 1.0, 1.0, 1.0], [2.0, 3.0, 7.0, 7.0]]]],
requires_grad=True)
self.assertONNX(nn.ReflectionPad2d((2, 3, 0, 1)), x)
def __init__(
self,
dim: int,
padding_type: str,
norm_layer: Type[nn.Module],
use_dropout: bool,
use_bias: bool,
n_conv_layers: int = 2,
dilations: List[int] = None,
):
"""Initialize the Resnet block.
A resnet block is a conv block with skip connections.
We implement skip connections in <forward> function. Original Resnet paper:
https://arxiv.org/pdf/1512.03385.pdf
Args:
dim: Number of channels in the conv layer
padding_type: Name of padding layer: reflect | replicate | zero
norm_layer: Normalization Layer class
use_dropout: Whether to use dropout layers
use_bias: Whether the conv layer uses bias or not
n_conv_layers: Number of convolution layers in this block
dilations: List of dilations for each conv layer
"""
super().__init__()
if dilations is None:
dilations = [1 for _ in range(n_conv_layers)]
assert n_conv_layers == len(
dilations
), "There must be exactly one dilation value for each conv layer."
conv_block = []
for dilation in dilations:
padding = 0
if padding_type == "reflect":
conv_block += [nn.ReflectionPad2d(dilation)]
elif padding_type == "replicate":
conv_block += [nn.ReplicationPad2d(dilation)]
elif padding_type == "zero":
padding = dilation
else:
raise NotImplementedError("padding [%s] is not implemented" %
padding_type)
conv_block += [
nn.Conv2d(
dim,
dim,
kernel_size=3,
padding=padding,
dilation=dilation,
bias=use_bias,
),
norm_layer(dim),
nn.ReLU(True),
]
if use_dropout:
conv_block += [nn.Dropout(0.5)]
if use_dropout:
# The last dropout layer should not be there
del conv_block[-1]
self.conv_block = nn.Sequential(*conv_block)
def __init__(self,
in_channels,
out_channels,
nBlocks,
nChanFirstConv=64,
dropout=False):
"""
nChanFirstConv: number of channels of the first convolution layer, 64 used in cycleGAN paper
"""
super(generator, self).__init__()
layers = []
layers += [nn.ReflectionPad2d(3)]
layers += [
nn.Conv2d(in_channels=in_channels,
out_channels=nChanFirstConv,
kernel_size=7)
]
layers += [nn.InstanceNorm2d(nChanFirstConv)]
layers += [nn.ReLU(inplace=True)]
layers += [
nn.Conv2d(in_channels=nChanFirstConv,
out_channels=nChanFirstConv * 2,
kernel_size=3,
stride=2,
padding=1)
]
layers += [nn.InstanceNorm2d(nChanFirstConv * 2)]
layers += [nn.ReLU(inplace=True)]
layers += [
nn.Conv2d(in_channels=nChanFirstConv * 2,
out_channels=nChanFirstConv * 4,
kernel_size=3,
stride=2,
padding=1)
]
layers += [nn.InstanceNorm2d(nChanFirstConv * 4)]
layers += [nn.ReLU(inplace=True)]
for i in range(nBlocks):
layers += [ResNetBlock(nChanFirstConv * 4, dropout)]
layers += [
nn.ConvTranspose2d(in_channels=nChanFirstConv * 4,
out_channels=nChanFirstConv * 2,
kernel_size=3,
stride=2,
padding=1,
output_padding=1)
]
layers += [nn.InstanceNorm2d(nChanFirstConv * 2)]
layers += [nn.ReLU(inplace=True)]
layers += [
nn.ConvTranspose2d(in_channels=nChanFirstConv * 2,
out_channels=nChanFirstConv,
kernel_size=3,
stride=2,
padding=1,
output_padding=1)
]
layers += [nn.InstanceNorm2d(nChanFirstConv)]
layers += [nn.ReLU(inplace=True)]
layers += [nn.ReflectionPad2d(3)]
layers += [
nn.Conv2d(in_channels=nChanFirstConv,
out_channels=out_channels,
kernel_size=7)
]
layers += [nn.Tanh()]
self.all_layers_ = nn.Sequential(*layers)
def __init__(self, input_nc=3, output_nc=3, ngf=64, use_dropout=False, n_blocks=9, padding_type='reflect', cfg=None, opt=None):
super(ResnetGenerator, self).__init__()
assert(n_blocks >= 0)
self.opt = opt
norm_layer = functools.partial(nn.InstanceNorm2d, affine=False, track_running_stats=False)
if type(norm_layer) == functools.partial:
use_bias = norm_layer.func == nn.InstanceNorm2d
else:
use_bias = norm_layer == nn.InstanceNorm2d
cfg_index = 0
model = [nn.ReflectionPad2d(3),
nn.Conv2d(input_nc, ngf if cfg is None else cfg[cfg_index], kernel_size=7, padding=0, bias=use_bias),
norm_layer(ngf),
nn.ReLU(True)]
cfg_index += 1
n_downsampling = 2
for i in range(n_downsampling):
mult = 2 ** i
input_channel_num = ngf * mult if cfg is None else cfg[cfg_index-1]
output_channel_num = ngf * mult * 2 if cfg is None else cfg[cfg_index]
cfg_index += 1
model += [nn.Conv2d(input_channel_num, output_channel_num, kernel_size=3, stride=2, padding=1, bias=use_bias),
norm_layer(output_channel_num),
nn.ReLU(True)]
mult = 2 ** n_downsampling
for i in range(n_blocks):
block_layer1_input_channel_num = ngf * mult if cfg is None else cfg[cfg_index-1]
block_layer1_output_channel_num = ngf * mult if cfg is None else cfg[cfg_index]
cfg_index += 1
block_layer2_output_channel_num = ngf * mult if cfg is None else cfg[cfg_index]
cfg_index += 1
if block_layer1_output_channel_num == 0:
continue
model += [ResnetBlock(block_layer1_input_channel_num,
block_layer1_output_channel_num,
block_layer2_output_channel_num,
padding_type=padding_type, norm_layer=norm_layer, use_dropout=use_dropout, use_bias=use_bias)]
output_channel_num = ngf
for i in range(n_downsampling):
mult = 2 ** (n_downsampling - i)
if i == n_downsampling - 1 and self.opt.unmask_last_upconv:
input_channel_num = ngf * mult if cfg is None or cfg_index == 0 else cfg[cfg_index - 1]
output_channel_num = int(ngf * mult / 2)
else:
input_channel_num = ngf * mult if cfg is None or cfg_index == 0 else cfg[cfg_index - 1]
output_channel_num = int(ngf * mult / 2) if cfg is None else cfg[cfg_index]
cfg_index += 1
model += [nn.ConvTranspose2d(input_channel_num, output_channel_num,
kernel_size=3, stride=2,
padding=1, output_padding=1,
bias=use_bias),
norm_layer(output_channel_num),
nn.ReLU(True)]
model += [nn.ReflectionPad2d(3)]
model += [nn.Conv2d(output_channel_num, output_nc, kernel_size=7, padding=0)]
model += [nn.Tanh()]
self.model = nn.Sequential(*model)
def __init__(self):
super(level_4_decoder, self).__init__()
d = load_t7("models/feature_invertor_conv4_1.t7")
# decoder
self.reflecPad11 = nn.ReflectionPad2d((1, 1, 1, 1))
self.conv11 = nn.Conv2d(512, 256, 3, 1, 0)
self.conv11.weight = torch.nn.Parameter(
torch.from_numpy(d[0]["w"]).cuda())
self.conv11.bias = torch.nn.Parameter(
torch.from_numpy(d[0]["b"]).cuda())
self.relu11 = nn.ReLU(inplace=True)
# 28 x 28
self.unpool = nn.UpsamplingNearest2d(scale_factor=2)
# 56 x 56
self.reflecPad12 = nn.ReflectionPad2d((1, 1, 1, 1))
self.conv12 = nn.Conv2d(256, 256, 3, 1, 0)
self.conv12.weight = torch.nn.Parameter(
torch.from_numpy(d[1]["w"]).cuda())
self.conv12.bias = torch.nn.Parameter(
torch.from_numpy(d[1]["b"]).cuda())
self.relu12 = nn.ReLU(inplace=True)
# 56 x 56
self.reflecPad13 = nn.ReflectionPad2d((1, 1, 1, 1))
self.conv13 = nn.Conv2d(256, 256, 3, 1, 0)
self.conv13.weight = torch.nn.Parameter(
torch.from_numpy(d[2]["w"]).cuda())
self.conv13.bias = torch.nn.Parameter(
torch.from_numpy(d[2]["b"]).cuda())
self.relu13 = nn.ReLU(inplace=True)
# 56 x 56
self.reflecPad14 = nn.ReflectionPad2d((1, 1, 1, 1))
self.conv14 = nn.Conv2d(256, 256, 3, 1, 0)
self.conv14.weight = torch.nn.Parameter(
torch.from_numpy(d[3]["w"]).cuda())
self.conv14.bias = torch.nn.Parameter(
torch.from_numpy(d[3]["b"]).cuda())
self.relu14 = nn.ReLU(inplace=True)
# 56 x 56
self.reflecPad15 = nn.ReflectionPad2d((1, 1, 1, 1))
self.conv15 = nn.Conv2d(256, 128, 3, 1, 0)
self.conv15.weight = torch.nn.Parameter(
torch.from_numpy(d[4]["w"]).cuda())
self.conv15.bias = torch.nn.Parameter(
torch.from_numpy(d[4]["b"]).cuda())
self.relu15 = nn.ReLU(inplace=True)
# 56 x 56
self.unpool2 = nn.UpsamplingNearest2d(scale_factor=2)
# 112 x 112
self.reflecPad16 = nn.ReflectionPad2d((1, 1, 1, 1))
self.conv16 = nn.Conv2d(128, 128, 3, 1, 0)
self.conv16.weight = torch.nn.Parameter(
torch.from_numpy(d[5]["w"]).cuda())
self.conv16.bias = torch.nn.Parameter(
torch.from_numpy(d[5]["b"]).cuda())
self.relu16 = nn.ReLU(inplace=True)
# 112 x 112
self.reflecPad17 = nn.ReflectionPad2d((1, 1, 1, 1))
self.conv17 = nn.Conv2d(128, 64, 3, 1, 0)
self.conv17.weight = torch.nn.Parameter(
torch.from_numpy(d[6]["w"]).cuda())
self.conv17.bias = torch.nn.Parameter(
torch.from_numpy(d[6]["b"]).cuda())
self.relu17 = nn.ReLU(inplace=True)
# 112 x 112
self.unpool3 = nn.UpsamplingNearest2d(scale_factor=2)
# 224 x 224
self.reflecPad18 = nn.ReflectionPad2d((1, 1, 1, 1))
self.conv18 = nn.Conv2d(64, 64, 3, 1, 0)
self.conv18.weight = torch.nn.Parameter(
torch.from_numpy(d[7]["w"]).cuda())
self.conv18.bias = torch.nn.Parameter(
torch.from_numpy(d[7]["b"]).cuda())
self.relu18 = nn.ReLU(inplace=True)
# 224 x 224
self.reflecPad19 = nn.ReflectionPad2d((1, 1, 1, 1))
self.conv19 = nn.Conv2d(64, 3, 3, 1, 0)
self.conv19.weight = torch.nn.Parameter(
torch.from_numpy(d[8]["w"]).cuda())
self.conv19.bias = torch.nn.Parameter(
torch.from_numpy(d[8]["b"]).cuda())
def __init__(self,
input_nc,
output_nc,
ngf=32,
norm_layer=nn.BatchNorm2d,
use_dropout=False,
n_blocks=3,
padding_type='reflect'):
super(Autoencoder, self).__init__()
use_bias = norm_layer == nn.InstanceNorm2d
model = [nn.Conv2d(input_nc, ngf, kernel_size=1)]
model = [ #nn.ReflectionPad2d(1),
nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0, bias=use_bias),
norm_layer(ngf),
nn.ReLU(True)
]
n_downsampling = 4
# Special case for 9th block of resnet
#n_downsampling, n_blocks = 0, 0
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),
norm_layer(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_layer=norm_layer,
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),
norm_layer(int(ngf * mult / 2)),
nn.ReLU(True)
]
n_upsampling_extra = 1
for i in range(n_upsampling_extra): # add upsampling layers
model += [
nn.ConvTranspose2d(ngf,
ngf,
kernel_size=3,
stride=2,
padding=1,
output_padding=1,
bias=use_bias),
norm_layer(ngf),
nn.ReLU(True)
]
if i == 3:
model += [
nn.Conv2d(ngf, ngf, kernel_size=3, stride=1, padding=0),
norm_layer(ngf),
nn.ReLU(True)
] #"""
model += [nn.ReflectionPad2d(3)]
model += [nn.Conv2d(ngf, ngf // 2, kernel_size=7, padding=0)]
model += [nn.ReflectionPad2d(3)]
model += [nn.Conv2d(ngf // 2, ngf // 4, kernel_size=5, padding=0)]
model += [nn.ReflectionPad2d(3)]
model += [nn.Conv2d(ngf // 4, output_nc, kernel_size=5, padding=0)]
model += [nn.ReflectionPad2d(3)]
model += [nn.Conv2d(output_nc, output_nc, kernel_size=5, padding=0)]
self.m = nn.Sequential(*model)
def __init__(self, input_nc, ndf=64, n_layers=5):
super().__init__()
model = [
nn.ReflectionPad2d(1),
nn.utils.spectral_norm(
nn.Conv2d(input_nc,
ndf,
kernel_size=4,
stride=2,
padding=0,
bias=True)),
nn.LeakyReLU(0.2, True)
]
for i in range(1, n_layers - 2):
mult = 2**(i - 1)
model += [
nn.ReflectionPad2d(1),
nn.utils.spectral_norm(
nn.Conv2d(ndf * mult,
ndf * mult * 2,
kernel_size=4,
stride=2,
padding=0,
bias=True)),
nn.LeakyReLU(0.2, True)
]
mult = 2**(n_layers - 2 - 1)
model += [
nn.ReflectionPad2d(1),
nn.utils.spectral_norm(
nn.Conv2d(ndf * mult,
ndf * mult * 2,
kernel_size=4,
stride=1,
padding=0,
bias=True)),
nn.LeakyReLU(0.2, True)
]
# Class Activation Map
mult = 2**(n_layers - 2)
self.gap_fc = nn.utils.spectral_norm(
nn.Linear(ndf * mult, 1, bias=False))
self.gmp_fc = nn.utils.spectral_norm(
nn.Linear(ndf * mult, 1, bias=False))
self.conv1x1 = nn.Conv2d(ndf * mult * 2,
ndf * mult,
kernel_size=1,
stride=1,
bias=True)
self.leaky_relu = nn.LeakyReLU(0.2, True)
self.pad = nn.ReflectionPad2d(1)
self.conv = nn.utils.spectral_norm(
nn.Conv2d(ndf * mult,
1,
kernel_size=4,
stride=1,
padding=0,
bias=False))
self.model = nn.Sequential(*model)
def forward(self, x):
norm = torch.norm(x, p=2, dim=1, keepdim=True)
clip = torch.clamp(norm, min=0)
expand = clip.expand_as(x)
return x / expand
model = NormL2()
x = Variable(torch.randn(1, 2, 3, 4))
save_data_and_model("reduceL2_subgraph", x, model)
model = nn.ZeroPad2d(1)
model.eval()
input = torch.rand(1, 3, 2, 4)
save_data_and_model("ZeroPad2d", input, model, version = 11)
model = nn.ReflectionPad2d(1)
model.eval()
input = torch.rand(1, 3, 2, 4)
save_data_and_model("ReflectionPad2d", input, model, version = 11)
# source: https://github.com/amdegroot/ssd.pytorch/blob/master/layers/modules/l2norm.py
class L2Norm(nn.Module):
def __init__(self,n_channels, scale):
super(L2Norm,self).__init__()
self.n_channels = n_channels
self.gamma = scale or None
self.eps = 1e-10
self.weight = nn.Parameter(torch.Tensor(self.n_channels))
self.reset_parameters()
def reset_parameters(self):
def __init__(self,
input_nc,
output_nc,
ngf=32,
n_downsample_global=3,
n_blocks_global=9,
n_local_enhancers=1,
n_blocks_local=3,
norm_layer=nn.BatchNorm2d,
padding_type='reflect'):
super(LocalEnhancer, self).__init__()
self.n_local_enhancers = n_local_enhancers
###### global generator model #####
ngf_global = ngf * (2**n_local_enhancers)
model_global = GlobalGenerator(input_nc, output_nc, ngf_global,
n_downsample_global, n_blocks_global,
norm_layer).model
model_global = [model_global[i] for i in range(len(model_global) - 3)
] # get rid of final convolution layers
self.model = nn.Sequential(*model_global)
###### local enhancer layers #####
for n in range(1, n_local_enhancers + 1):
### downsample
ngf_global = ngf * (2**(n_local_enhancers - n))
model_downsample = [
nn.ReflectionPad2d(3),
nn.Conv2d(input_nc, ngf_global, kernel_size=7, padding=0),
norm_layer(ngf_global),
nn.ReLU(True),
nn.Conv2d(ngf_global,
ngf_global * 2,
kernel_size=3,
stride=2,
padding=1),
norm_layer(ngf_global * 2),
nn.ReLU(True)
]
### residual blocks
model_upsample = []
for i in range(n_blocks_local):
model_upsample += [
ResnetBlock(ngf_global * 2,
padding_type=padding_type,
norm_layer=norm_layer)
]
### upsample
model_upsample += [
nn.ConvTranspose2d(ngf_global * 2,
ngf_global,
kernel_size=3,
stride=2,
padding=1,
output_padding=1),
norm_layer(ngf_global),
nn.ReLU(True)
]
### final convolution
if n == n_local_enhancers:
model_upsample += [
nn.ReflectionPad2d(3),
nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0),
nn.Tanh()
]
setattr(self, 'model' + str(n) + '_1',
nn.Sequential(*model_downsample))
setattr(self, 'model' + str(n) + '_2',
nn.Sequential(*model_upsample))
self.downsample = nn.AvgPool2d(3,
stride=2,
padding=[1, 1],
count_include_pad=False)
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,
input_nc,
output_nc,
ngf=64,
n_blocks=6,
norm='batch',
activation='PReLU',
drop_rate=0,
add_noise=False,
gpu_ids=[]):
super(_ResGenerator, self).__init__()
self.gpu_ids = gpu_ids
norm_layer = get_norm_layer(norm_type=norm)
nonlinearity = get_nonlinearity_layer(activation_type=activation)
if type(norm_layer) == functools.partial:
use_bias = norm_layer.func == nn.InstanceNorm2d
else:
use_bias = norm_layer == nn.InstanceNorm2d
encoder = [
nn.ReflectionPad2d(3),
nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0, bias=use_bias),
norm_layer(ngf), nonlinearity
]
n_downsampling = 2
mult = 1
for i in range(n_downsampling):
mult_prev = mult
mult = min(2**(i + 1), 2)
encoder += [
_EncoderBlock(ngf * mult_prev, ngf * mult, ngf * mult,
norm_layer, nonlinearity, use_bias),
nn.AvgPool2d(kernel_size=2, stride=2)
]
mult = min(2**n_downsampling, 2)
for i in range(n_blocks - n_downsampling):
encoder += [
_InceptionBlock(ngf * mult,
ngf * mult,
norm_layer=norm_layer,
nonlinearity=nonlinearity,
width=1,
drop_rate=drop_rate,
use_bias=use_bias)
]
decoder = []
if add_noise:
decoder += [GaussianNoiseLayer()]
for i in range(n_downsampling):
mult_prev = mult
mult = min(2**(n_downsampling - i - 1), 2)
decoder += [
_DecoderUpBlock(ngf * mult_prev, ngf * mult_prev, ngf * mult,
norm_layer, nonlinearity, use_bias),
]
decoder += [
nn.ReflectionPad2d(3),
nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0),
nn.Tanh()
]
self.encoder = nn.Sequential(*encoder)
self.decoder = nn.Sequential(*decoder)
def __init__(self,
input_nc,
output_nc,
ngf=64,
n_downsampling=3,
n_blocks=9,
norm_layer=nn.BatchNorm2d,
padding_type='reflect'):
assert (n_blocks >= 0)
super(GlobalGenerator, self).__init__()
activation = nn.ReLU(True)
# model_down_1 = [
# # nn.ReflectionPad2d(3),
# nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0)
#
# ]
model_down = [
nn.ReflectionPad2d(3),
nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0),
norm_layer(ngf), activation
]
### downsample
for i in range(n_downsampling):
mult = 2**i
model_down += [
nn.Conv2d(ngf * mult,
ngf * mult * 2,
kernel_size=3,
stride=2,
padding=1),
norm_layer(ngf * mult * 2), activation
]
# self.model_down_1 = nn.Sequential(*model_down_1)
self.model_down = nn.Sequential(*model_down)
### resnet blocks
mult = 2**n_downsampling
model_res = []
for i in range(n_blocks):
model_res += [
ResnetBlock(ngf * mult,
padding_type=padding_type,
activation=activation,
norm_layer=norm_layer)
]
self.model_res = nn.Sequential(*model_res)
### upsample
model_up = []
for i in range(n_downsampling):
mult = 2**(n_downsampling - i)
model_up += [
nn.ConvTranspose2d(ngf * mult,
int(ngf * mult / 2),
kernel_size=3,
stride=2,
padding=1,
output_padding=1),
norm_layer(int(ngf * mult / 2)), activation
]
model_up += [
nn.ReflectionPad2d(3),
nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0),
nn.Tanh()
]
self.model = nn.Sequential(*model_up)
def __init__(self, in_channels, out_channels, kernel_size, stride):
super(ConvLayer, self).__init__()
reflection_padding = int(np.floor(kernel_size / 2))
self.reflection_pad = nn.ReflectionPad2d(reflection_padding)
self.conv2d = nn.Conv2d(in_channels, out_channels, kernel_size, stride)
def __init__(self,
in_channels,
out_channels,
conv_type,
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:
print("sn")
self.conv2d = SpectralNorm(
nn.Conv2d(in_channels,
out_channels,
kernel_size,
stride,
padding=0,
dilation=dilation))
else:
if conv_type == 'normal':
self.conv2d = nn.Conv2d(in_channels,
out_channels,
kernel_size,
stride,
padding=0,
dilation=dilation)
elif conv_type == 'partial':
self.conv2d = PartialConv2d(in_channels,
out_channels,
kernel_size,
stride,
padding=0,
dilation=dilation)
else:
print("conv_type not implemented")
def __init__(self,
input_nc,
output_nc,
ngf=64,
norm_layer=nn.BatchNorm2d,
use_dropout=False,
n_blocks=6,
gpu_ids=[],
padding_type='reflect'):
assert (n_blocks >= 0)
super(ResnetGenerator, self).__init__()
self.input_nc = input_nc
self.output_nc = output_nc
self.ngf = ngf
self.gpu_ids = gpu_ids
if type(norm_layer) == functools.partial:
use_bias = norm_layer.func == nn.InstanceNorm2d
else:
use_bias = norm_layer == nn.InstanceNorm2d
model = [
nn.ReflectionPad2d(3),
nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0, bias=use_bias),
norm_layer(ngf),
nn.ReLU(True)
]
n_downsampling = 2
for i in range(n_downsampling):
mult = 2**i
model += [
nn.Conv2d(ngf * mult,
ngf * mult * 2,
kernel_size=3,
stride=2,
padding=1,
bias=use_bias),
norm_layer(ngf * mult * 2),
nn.ReLU(True)
]
mult = 2**n_downsampling
for i in range(n_blocks):
model += [
ResnetBlock(ngf * mult,
padding_type=padding_type,
norm_layer=norm_layer,
use_dropout=use_dropout,
use_bias=use_bias)
]
for i in range(n_downsampling):
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),
norm_layer(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, in_channels, out_channels,
extra_channels,
ngf=32,
norm_layer=nn.InstanceNorm2d,
use_dropout=False,
n_blocks=6,
padding_type='reflect'):
assert(n_blocks >= 0)
super(StretchNet, self).__init__()
self.extra_channels = extra_channels
self.in_channels = in_channels
self.out_channels = out_channels
self.ngf = ngf
if type(norm_layer) == functools.partial:
use_bias = norm_layer.func == nn.InstanceNorm2d
else:
use_bias = norm_layer == nn.InstanceNorm2d
models = []
models.append(
(nn.Sequential(
nn.ReflectionPad2d(3),
nn.Conv2d(
in_channels+extra_channels,
ngf, kernel_size=7, padding=0,
bias=use_bias),
norm_layer(ngf),
nn.ReLU(True)
), 1)
)
n_downsampling = 1
for i in range(n_downsampling):
mult = 4**i
models.append(
(nn.Sequential(
nn.Conv2d(
ngf * mult + extra_channels,
ngf * mult * 4, kernel_size=3,
stride=2, padding=1, bias=use_bias
),
norm_layer(ngf * mult * 4),
nn.ReLU(True)
), 1)
)
mult = 4**n_downsampling
for i in range(n_blocks):
models.append(
(
StretchBlock(
ngf * mult, 64//(2**n_downsampling), 64//(2**n_downsampling),
self.extra_channels,
),0)
)
pass
for i in range(n_downsampling):
mult = 4**(n_downsampling - i)
models.append(
(nn.Sequential(
nn.ConvTranspose2d(ngf * mult + extra_channels,
int(ngf * mult / 4),
kernel_size=3, stride=2,
padding=1, output_padding=1,
bias=use_bias),
norm_layer(int(ngf * mult / 4)),
nn.ReLU(True)
), 1)
)
models.append(
(
nn.Sequential(
nn.ReflectionPad2d(3),
nn.Conv2d(ngf, out_channels, kernel_size=7, padding=0),
nn.Tanh()
), 2
)
)
self.models = models
for i, (model,tag) in enumerate(self.models):
setattr(self, 'model_'+str(i), model)