def define_D(input_nc,
ndf,
netD,
n_layers_D=3,
norm='batch',
init_type='normal',
init_gain=0.02,
gpu_ids=[]):
"""Create a discriminator
Parameters:
input_nc (int) -- the number of channels in input images
ndf (int) -- the number of filters in the first conv layer
netD (str) -- the architecture's name: basic | n_layers | pixel
n_layers_D (int) -- the number of conv layers in the discriminator; effective when netD=='n_layers'
norm (str) -- the type of normalization layers used in the network.
init_type (str) -- the name of the initialization method.
init_gain (float) -- scaling factor for normal, xavier and orthogonal.
gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2
Returns a discriminator
Our current implementation provides three types of discriminators:
[basic]: 'PatchGAN' classifier described in the original pix2pix paper.
It can classify whether 70×70 overlapping patches are real or fake.
Such a patch-level discriminator architecture has fewer parameters
than a full-image discriminator and can work on arbitrarily-sized images
in a fully convolutional fashion.
[n_layers]: With this mode, you cna specify the number of conv layers in the discriminator
with the parameter <n_layers_D> (default=3 as used in [basic] (PatchGAN).)
[pixel]: 1x1 PixelGAN discriminator can classify whether a pixel is real or not.
It encourages greater color diversity but has no effect on spatial statistics.
The discriminator has been initialized by <init_net>. It uses Leakly RELU for non-linearity.
"""
net = None
norm_layer = utils.get_norm_layer(norm_type=norm)
if netD == 'basic': # default PatchGAN classifier
net = NLayerDiscriminator(input_nc,
ndf,
n_layers=3,
norm_layer=norm_layer)
elif netD == 'n_layers': # more options
net = NLayerDiscriminator(input_nc,
ndf,
n_layers_D,
norm_layer=norm_layer)
elif netD == 'pixel': # classify if each pixel is real or fake
net = PixelDiscriminator(input_nc, ndf, norm_layer=norm_layer)
else:
raise NotImplementedError(
'Discriminator models name [%s] is not recognized' % netD)
return init_net(net, init_type, init_gain, gpu_ids)
def __init__(self,
in_dims=5,
out_dims=1,
n_res=5,
nf=32,
norm_type='instance'):
super(MultiResnet, self).__init__()
norm_layer = utils.get_norm_layer(norm_type=norm_type)
self.x3_conv = DoubleConv(in_dims,
nf,
kernel_size=3,
stride=1,
padding=1,
use_bias=False,
norm_layer=norm_layer)
self.x5_conv = DoubleConv(in_dims,
nf,
kernel_size=5,
stride=1,
padding=2,
use_bias=False,
norm_layer=norm_layer)
self.x7_conv = DoubleConv(in_dims,
nf,
kernel_size=7,
stride=1,
padding=3,
use_bias=False,
norm_layer=norm_layer)
features = []
for i in range(n_res):
features += [
ResnetBlock(nf,
kernel_size=3,
stride=1,
padding=1,
use_bias=False,
norm_layer=norm_layer)
]
self.features = nn.Sequential(*features)
self.final_conv = nn.Conv2d(nf * 3,
out_dims,
3,
1,
padding=1,
bias=False)
def __init__(self,
input_dims,
nf=32,
n_layers=4,
norm_type='instance',
use_bias=False):
super(NAttentionDiscriminator, self).__init__()
norm_layer = utils.get_norm_layer(norm_type=norm_type)
layers = [
nn.utils.spectral_norm(
nn.Conv2d(input_dims, nf, 3, 1, padding=1, bias=False)),
# norm_layer(nf),
nn.ReLU(True),
nn.utils.spectral_norm(
nn.Conv2d(nf, nf, 3, 1, padding=1, bias=False)),
# norm_layer(nf),
nn.ReLU(True)
]
# for i in range(2):
layers += [
ResAttModule(nf,
kernel_size=3,
stride=1,
padding=1,
use_bias=False,
norm_layer=norm_layer)
]
for n in range(1, n_layers):
layers += [
nn.utils.spectral_norm(
nn.Conv2d(nf,
nf,
kernel_size=3,
stride=2,
padding=0,
bias=use_bias)),
# norm_layer(nf),
# nn.LeakyReLU(0.2, True)
nn.ReLU(True)
]
layers += [nn.Conv2d(nf, 1, kernel_size=3, stride=2)]
self.model = nn.Sequential(*layers)
def define_G(input_nc,
output_nc,
ngf,
netG,
norm='batch',
use_dropout=False,
init_type='normal',
init_gain=0.02,
gpu_ids=[]):
"""Create a 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
netG (str) -- the architecture's name: resnet_9blocks | resnet_6blocks | unet_256 | unet_128
norm (str) -- the name of normalization layers used in the network: batch | instance | none
use_dropout (bool) -- if use dropout layers.
init_type (str) -- the name of our initialization method.
init_gain (float) -- scaling factor for normal, xavier and orthogonal.
gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2
Returns a generator
Our current implementation provides two types of generators:
U-Net: [unet_128] (for 128x128 input images) and [unet_256] (for 256x256 input images)
The original U-Net paper: https://arxiv.org/abs/1505.04597
Resnet-based generator: [resnet_6blocks] (with 6 Resnet blocks) and [resnet_9blocks] (with 9 Resnet blocks)
Resnet-based generator consists of several Resnet blocks between a few downsampling/upsampling operations.
We adapt Torch code from Justin Johnson's neural style transfer project (https://github.com/jcjohnson/fast-neural-style).
The generator has been initialized by <init_net>. It uses RELU for non-linearity.
"""
net = None
norm_layer = utils.get_norm_layer(norm_type=norm)
if netG == 'resnet_9blocks':
net = ResnetGenerator(input_nc,
output_nc,
ngf,
norm_layer=norm_layer,
use_dropout=use_dropout,
n_blocks=9)
elif netG == 'resnet_6blocks':
net = ResnetGenerator(input_nc,
output_nc,
ngf,
norm_layer=norm_layer,
use_dropout=use_dropout,
n_blocks=6)
elif netG == 'unet_128':
net = UnetGenerator(input_nc,
output_nc,
7,
ngf,
norm_layer=norm_layer,
use_dropout=use_dropout)
elif netG == 'unet_64':
net = UnetGenerator(input_nc,
output_nc,
5,
ngf,
norm_layer=norm_layer,
use_dropout=use_dropout)
elif netG == 'unet_256':
net = UnetGenerator(input_nc,
output_nc,
8,
ngf,
norm_layer=norm_layer,
use_dropout=use_dropout)
else:
raise NotImplementedError(
'Generator models name [%s] is not recognized' % netG)
return init_net(net, init_type, init_gain, gpu_ids)
def __init__(self,
in_dims=5,
out_dims=1,
n_res=6,
n_att=2,
nf=32,
norm_type='instance'):
super(MultiAttResnet, self).__init__()
norm_layer = utils.get_norm_layer(norm_type=norm_type)
norm = True
use_bias = False
# ------> 浅层特征 <------ #
self.base_block = FeatureModule(in_dims,
nf=nf,
use_bias=use_bias,
norm=norm,
norm_layer=norm_layer)
# ------> 注意力特征 <------ #
self.att_blocks = nn.Sequential(
# 下采样
nn.ReflectionPad2d(1),
nn.Conv2d(nf, nf * 2, kernel_size=3, stride=2),
# 注意力模块
ResAttModule(nf * 2,
kernel_size=3,
stride=1,
padding=1,
use_bias=use_bias,
norm=norm,
norm_layer=norm_layer),
# 下采样
nn.ReflectionPad2d(1),
nn.Conv2d(nf * 2, nf * 2, kernel_size=3, stride=2),
# 注意力模块
ResAttModule(nf * 2,
kernel_size=3,
stride=1,
padding=1,
use_bias=use_bias,
norm=norm,
norm_layer=norm_layer),
# 上采样
nn.ConvTranspose2d(nf * 2,
nf * 2,
3,
stride=2,
padding=1,
output_padding=1),
# 注意力模块
ResAttModule(nf * 2,
kernel_size=3,
stride=1,
padding=1,
use_bias=use_bias,
norm=norm,
norm_layer=norm_layer),
# 上采样
nn.ConvTranspose2d(nf * 2,
nf,
3,
stride=2,
padding=1,
output_padding=1))
# ------> 主干特征 <------ #
# 两次下采样
self.res_blocks = nn.Sequential(
# 下采样
nn.ReflectionPad2d(1),
nn.Conv2d(nf, nf * 2, kernel_size=3, stride=2),
# 两层残差
ResnetBlock(nf * 2,
kernel_size=3,
stride=1,
padding=1,
use_bias=use_bias,
norm=norm,
norm_layer=norm_layer),
ResnetBlock(nf * 2,
kernel_size=3,
stride=1,
padding=1,
use_bias=use_bias,
norm=norm,
norm_layer=norm_layer),
# 下采样
nn.ReflectionPad2d(1),
nn.Conv2d(nf * 2, nf * 2, kernel_size=3, stride=2),
# 2层残差
ResnetBlock(nf * 2,
kernel_size=3,
stride=1,
padding=1,
use_bias=use_bias,
norm=norm,
norm_layer=norm_layer),
ResnetBlock(nf * 2,
kernel_size=3,
stride=1,
padding=1,
use_bias=use_bias,
norm=norm,
norm_layer=norm_layer),
# 上采样
nn.ConvTranspose2d(nf * 2,
nf * 2,
3,
stride=2,
padding=1,
output_padding=1),
# 两层残差
ResnetBlock(nf * 2,
kernel_size=3,
stride=1,
padding=1,
use_bias=use_bias,
norm=norm,
norm_layer=norm_layer),
ResnetBlock(nf * 2,
kernel_size=3,
stride=1,
padding=1,
use_bias=use_bias,
norm=norm,
norm_layer=norm_layer),
# 上采样
nn.ConvTranspose2d(nf * 2,
nf,
3,
stride=2,
padding=1,
output_padding=1))
# ------> 输出层 <------ #
self.out_blocks = nn.Sequential(
# DoubleConv(nf * 3, nf, kernel_size=3, stride=1, padding=1, use_bias=use_bias, norm=norm,
# norm_layer=norm_layer),
nn.Conv2d(nf * 3, nf, 3, 1, padding=1, bias=use_bias),
# ResnetBlock(nf, kernel_size=3, stride=1, padding=1, use_bias=use_bias, norm=norm, norm_layer=norm_layer),
# ResnetBlock(nf, kernel_size=3, stride=1, padding=1, use_bias=use_bias, norm=norm, norm_layer=norm_layer),
nn.ReLU(inplace=True),
nn.Conv2d(nf, out_dims, 3, 1, padding=1, bias=use_bias))