368 lines
14 KiB
Python
368 lines
14 KiB
Python
import sys
|
|
import math
|
|
|
|
import torch
|
|
from torch import nn
|
|
from torch.nn import functional as F
|
|
from torch.autograd import Function
|
|
from torch.nn.utils.spectral_norm import spectral_norm as SpectralNorm
|
|
|
|
|
|
class LayerNorm2d(nn.Module):
|
|
def __init__(self, n_out, affine=True):
|
|
super(LayerNorm2d, self).__init__()
|
|
self.n_out = n_out
|
|
self.affine = affine
|
|
|
|
if self.affine:
|
|
self.weight = nn.Parameter(torch.ones(n_out, 1, 1))
|
|
self.bias = nn.Parameter(torch.zeros(n_out, 1, 1))
|
|
|
|
def forward(self, x):
|
|
normalized_shape = x.size()[1:]
|
|
if self.affine:
|
|
return F.layer_norm(x, normalized_shape, \
|
|
self.weight.expand(normalized_shape),
|
|
self.bias.expand(normalized_shape))
|
|
|
|
else:
|
|
return F.layer_norm(x, normalized_shape)
|
|
|
|
class ADAINHourglass(nn.Module):
|
|
def __init__(self, image_nc, pose_nc, ngf, img_f, encoder_layers, decoder_layers, nonlinearity, use_spect):
|
|
super(ADAINHourglass, self).__init__()
|
|
self.encoder = ADAINEncoder(image_nc, pose_nc, ngf, img_f, encoder_layers, nonlinearity, use_spect)
|
|
self.decoder = ADAINDecoder(pose_nc, ngf, img_f, encoder_layers, decoder_layers, True, nonlinearity, use_spect)
|
|
self.output_nc = self.decoder.output_nc
|
|
|
|
def forward(self, x, z):
|
|
return self.decoder(self.encoder(x, z), z)
|
|
|
|
|
|
|
|
class ADAINEncoder(nn.Module):
|
|
def __init__(self, image_nc, pose_nc, ngf, img_f, layers, nonlinearity=nn.LeakyReLU(), use_spect=False):
|
|
super(ADAINEncoder, self).__init__()
|
|
self.layers = layers
|
|
self.input_layer = nn.Conv2d(image_nc, ngf, kernel_size=7, stride=1, padding=3)
|
|
for i in range(layers):
|
|
in_channels = min(ngf * (2**i), img_f)
|
|
out_channels = min(ngf *(2**(i+1)), img_f)
|
|
model = ADAINEncoderBlock(in_channels, out_channels, pose_nc, nonlinearity, use_spect)
|
|
setattr(self, 'encoder' + str(i), model)
|
|
self.output_nc = out_channels
|
|
|
|
def forward(self, x, z):
|
|
out = self.input_layer(x)
|
|
out_list = [out]
|
|
for i in range(self.layers):
|
|
model = getattr(self, 'encoder' + str(i))
|
|
out = model(out, z)
|
|
out_list.append(out)
|
|
return out_list
|
|
|
|
class ADAINDecoder(nn.Module):
|
|
"""docstring for ADAINDecoder"""
|
|
def __init__(self, pose_nc, ngf, img_f, encoder_layers, decoder_layers, skip_connect=True,
|
|
nonlinearity=nn.LeakyReLU(), use_spect=False):
|
|
|
|
super(ADAINDecoder, self).__init__()
|
|
self.encoder_layers = encoder_layers
|
|
self.decoder_layers = decoder_layers
|
|
self.skip_connect = skip_connect
|
|
use_transpose = True
|
|
|
|
for i in range(encoder_layers-decoder_layers, encoder_layers)[::-1]:
|
|
in_channels = min(ngf * (2**(i+1)), img_f)
|
|
in_channels = in_channels*2 if i != (encoder_layers-1) and self.skip_connect else in_channels
|
|
out_channels = min(ngf * (2**i), img_f)
|
|
model = ADAINDecoderBlock(in_channels, out_channels, out_channels, pose_nc, use_transpose, nonlinearity, use_spect)
|
|
setattr(self, 'decoder' + str(i), model)
|
|
|
|
self.output_nc = out_channels*2 if self.skip_connect else out_channels
|
|
|
|
def forward(self, x, z):
|
|
out = x.pop() if self.skip_connect else x
|
|
for i in range(self.encoder_layers-self.decoder_layers, self.encoder_layers)[::-1]:
|
|
model = getattr(self, 'decoder' + str(i))
|
|
out = model(out, z)
|
|
out = torch.cat([out, x.pop()], 1) if self.skip_connect else out
|
|
return out
|
|
|
|
class ADAINEncoderBlock(nn.Module):
|
|
def __init__(self, input_nc, output_nc, feature_nc, nonlinearity=nn.LeakyReLU(), use_spect=False):
|
|
super(ADAINEncoderBlock, self).__init__()
|
|
kwargs_down = {'kernel_size': 4, 'stride': 2, 'padding': 1}
|
|
kwargs_fine = {'kernel_size': 3, 'stride': 1, 'padding': 1}
|
|
|
|
self.conv_0 = spectral_norm(nn.Conv2d(input_nc, output_nc, **kwargs_down), use_spect)
|
|
self.conv_1 = spectral_norm(nn.Conv2d(output_nc, output_nc, **kwargs_fine), use_spect)
|
|
|
|
|
|
self.norm_0 = ADAIN(input_nc, feature_nc)
|
|
self.norm_1 = ADAIN(output_nc, feature_nc)
|
|
self.actvn = nonlinearity
|
|
|
|
def forward(self, x, z):
|
|
x = self.conv_0(self.actvn(self.norm_0(x, z)))
|
|
x = self.conv_1(self.actvn(self.norm_1(x, z)))
|
|
return x
|
|
|
|
class ADAINDecoderBlock(nn.Module):
|
|
def __init__(self, input_nc, output_nc, hidden_nc, feature_nc, use_transpose=True, nonlinearity=nn.LeakyReLU(), use_spect=False):
|
|
super(ADAINDecoderBlock, self).__init__()
|
|
# Attributes
|
|
self.actvn = nonlinearity
|
|
hidden_nc = min(input_nc, output_nc) if hidden_nc is None else hidden_nc
|
|
|
|
kwargs_fine = {'kernel_size':3, 'stride':1, 'padding':1}
|
|
if use_transpose:
|
|
kwargs_up = {'kernel_size':3, 'stride':2, 'padding':1, 'output_padding':1}
|
|
else:
|
|
kwargs_up = {'kernel_size':3, 'stride':1, 'padding':1}
|
|
|
|
# create conv layers
|
|
self.conv_0 = spectral_norm(nn.Conv2d(input_nc, hidden_nc, **kwargs_fine), use_spect)
|
|
if use_transpose:
|
|
self.conv_1 = spectral_norm(nn.ConvTranspose2d(hidden_nc, output_nc, **kwargs_up), use_spect)
|
|
self.conv_s = spectral_norm(nn.ConvTranspose2d(input_nc, output_nc, **kwargs_up), use_spect)
|
|
else:
|
|
self.conv_1 = nn.Sequential(spectral_norm(nn.Conv2d(hidden_nc, output_nc, **kwargs_up), use_spect),
|
|
nn.Upsample(scale_factor=2))
|
|
self.conv_s = nn.Sequential(spectral_norm(nn.Conv2d(input_nc, output_nc, **kwargs_up), use_spect),
|
|
nn.Upsample(scale_factor=2))
|
|
# define normalization layers
|
|
self.norm_0 = ADAIN(input_nc, feature_nc)
|
|
self.norm_1 = ADAIN(hidden_nc, feature_nc)
|
|
self.norm_s = ADAIN(input_nc, feature_nc)
|
|
|
|
def forward(self, x, z):
|
|
x_s = self.shortcut(x, z)
|
|
dx = self.conv_0(self.actvn(self.norm_0(x, z)))
|
|
dx = self.conv_1(self.actvn(self.norm_1(dx, z)))
|
|
out = x_s + dx
|
|
return out
|
|
|
|
def shortcut(self, x, z):
|
|
x_s = self.conv_s(self.actvn(self.norm_s(x, z)))
|
|
return x_s
|
|
|
|
|
|
def spectral_norm(module, use_spect=True):
|
|
"""use spectral normal layer to stable the training process"""
|
|
if use_spect:
|
|
return SpectralNorm(module)
|
|
else:
|
|
return module
|
|
|
|
|
|
class ADAIN(nn.Module):
|
|
def __init__(self, norm_nc, feature_nc):
|
|
super().__init__()
|
|
|
|
self.param_free_norm = nn.InstanceNorm2d(norm_nc, affine=False)
|
|
|
|
nhidden = 128
|
|
use_bias=True
|
|
|
|
self.mlp_shared = nn.Sequential(
|
|
nn.Linear(feature_nc, nhidden, bias=use_bias),
|
|
nn.ReLU()
|
|
)
|
|
self.mlp_gamma = nn.Linear(nhidden, norm_nc, bias=use_bias)
|
|
self.mlp_beta = nn.Linear(nhidden, norm_nc, bias=use_bias)
|
|
|
|
def forward(self, x, feature):
|
|
|
|
# Part 1. generate parameter-free normalized activations
|
|
normalized = self.param_free_norm(x)
|
|
|
|
# Part 2. produce scaling and bias conditioned on feature
|
|
feature = feature.view(feature.size(0), -1)
|
|
actv = self.mlp_shared(feature)
|
|
gamma = self.mlp_gamma(actv)
|
|
beta = self.mlp_beta(actv)
|
|
|
|
# apply scale and bias
|
|
gamma = gamma.view(*gamma.size()[:2], 1,1)
|
|
beta = beta.view(*beta.size()[:2], 1,1)
|
|
out = normalized * (1 + gamma) + beta
|
|
return out
|
|
|
|
|
|
class FineEncoder(nn.Module):
|
|
"""docstring for Encoder"""
|
|
def __init__(self, image_nc, ngf, img_f, layers, norm_layer=nn.BatchNorm2d, nonlinearity=nn.LeakyReLU(), use_spect=False):
|
|
super(FineEncoder, self).__init__()
|
|
self.layers = layers
|
|
self.first = FirstBlock2d(image_nc, ngf, norm_layer, nonlinearity, use_spect)
|
|
for i in range(layers):
|
|
in_channels = min(ngf*(2**i), img_f)
|
|
out_channels = min(ngf*(2**(i+1)), img_f)
|
|
model = DownBlock2d(in_channels, out_channels, norm_layer, nonlinearity, use_spect)
|
|
setattr(self, 'down' + str(i), model)
|
|
self.output_nc = out_channels
|
|
|
|
def forward(self, x):
|
|
x = self.first(x)
|
|
out=[x]
|
|
for i in range(self.layers):
|
|
model = getattr(self, 'down'+str(i))
|
|
x = model(x)
|
|
out.append(x)
|
|
return out
|
|
|
|
class FineDecoder(nn.Module):
|
|
"""docstring for FineDecoder"""
|
|
def __init__(self, image_nc, feature_nc, ngf, img_f, layers, num_block, norm_layer=nn.BatchNorm2d, nonlinearity=nn.LeakyReLU(), use_spect=False):
|
|
super(FineDecoder, self).__init__()
|
|
self.layers = layers
|
|
for i in range(layers)[::-1]:
|
|
in_channels = min(ngf*(2**(i+1)), img_f)
|
|
out_channels = min(ngf*(2**i), img_f)
|
|
up = UpBlock2d(in_channels, out_channels, norm_layer, nonlinearity, use_spect)
|
|
res = FineADAINResBlocks(num_block, in_channels, feature_nc, norm_layer, nonlinearity, use_spect)
|
|
jump = Jump(out_channels, norm_layer, nonlinearity, use_spect)
|
|
|
|
setattr(self, 'up' + str(i), up)
|
|
setattr(self, 'res' + str(i), res)
|
|
setattr(self, 'jump' + str(i), jump)
|
|
|
|
self.final = FinalBlock2d(out_channels, image_nc, use_spect, 'tanh')
|
|
|
|
self.output_nc = out_channels
|
|
|
|
def forward(self, x, z):
|
|
out = x.pop()
|
|
for i in range(self.layers)[::-1]:
|
|
res_model = getattr(self, 'res' + str(i))
|
|
up_model = getattr(self, 'up' + str(i))
|
|
jump_model = getattr(self, 'jump' + str(i))
|
|
out = res_model(out, z)
|
|
out = up_model(out)
|
|
out = jump_model(x.pop()) + out
|
|
out_image = self.final(out)
|
|
return out_image
|
|
|
|
class FirstBlock2d(nn.Module):
|
|
"""
|
|
Downsampling block for use in encoder.
|
|
"""
|
|
def __init__(self, input_nc, output_nc, norm_layer=nn.BatchNorm2d, nonlinearity=nn.LeakyReLU(), use_spect=False):
|
|
super(FirstBlock2d, self).__init__()
|
|
kwargs = {'kernel_size': 7, 'stride': 1, 'padding': 3}
|
|
conv = spectral_norm(nn.Conv2d(input_nc, output_nc, **kwargs), use_spect)
|
|
|
|
if type(norm_layer) == type(None):
|
|
self.model = nn.Sequential(conv, nonlinearity)
|
|
else:
|
|
self.model = nn.Sequential(conv, norm_layer(output_nc), nonlinearity)
|
|
|
|
|
|
def forward(self, x):
|
|
out = self.model(x)
|
|
return out
|
|
|
|
class DownBlock2d(nn.Module):
|
|
def __init__(self, input_nc, output_nc, norm_layer=nn.BatchNorm2d, nonlinearity=nn.LeakyReLU(), use_spect=False):
|
|
super(DownBlock2d, self).__init__()
|
|
|
|
|
|
kwargs = {'kernel_size': 3, 'stride': 1, 'padding': 1}
|
|
conv = spectral_norm(nn.Conv2d(input_nc, output_nc, **kwargs), use_spect)
|
|
pool = nn.AvgPool2d(kernel_size=(2, 2))
|
|
|
|
if type(norm_layer) == type(None):
|
|
self.model = nn.Sequential(conv, nonlinearity, pool)
|
|
else:
|
|
self.model = nn.Sequential(conv, norm_layer(output_nc), nonlinearity, pool)
|
|
|
|
def forward(self, x):
|
|
out = self.model(x)
|
|
return out
|
|
|
|
class UpBlock2d(nn.Module):
|
|
def __init__(self, input_nc, output_nc, norm_layer=nn.BatchNorm2d, nonlinearity=nn.LeakyReLU(), use_spect=False):
|
|
super(UpBlock2d, self).__init__()
|
|
kwargs = {'kernel_size': 3, 'stride': 1, 'padding': 1}
|
|
conv = spectral_norm(nn.Conv2d(input_nc, output_nc, **kwargs), use_spect)
|
|
if type(norm_layer) == type(None):
|
|
self.model = nn.Sequential(conv, nonlinearity)
|
|
else:
|
|
self.model = nn.Sequential(conv, norm_layer(output_nc), nonlinearity)
|
|
|
|
def forward(self, x):
|
|
out = self.model(F.interpolate(x, scale_factor=2))
|
|
return out
|
|
|
|
class FineADAINResBlocks(nn.Module):
|
|
def __init__(self, num_block, input_nc, feature_nc, norm_layer=nn.BatchNorm2d, nonlinearity=nn.LeakyReLU(), use_spect=False):
|
|
super(FineADAINResBlocks, self).__init__()
|
|
self.num_block = num_block
|
|
for i in range(num_block):
|
|
model = FineADAINResBlock2d(input_nc, feature_nc, norm_layer, nonlinearity, use_spect)
|
|
setattr(self, 'res'+str(i), model)
|
|
|
|
def forward(self, x, z):
|
|
for i in range(self.num_block):
|
|
model = getattr(self, 'res'+str(i))
|
|
x = model(x, z)
|
|
return x
|
|
|
|
class Jump(nn.Module):
|
|
def __init__(self, input_nc, norm_layer=nn.BatchNorm2d, nonlinearity=nn.LeakyReLU(), use_spect=False):
|
|
super(Jump, self).__init__()
|
|
kwargs = {'kernel_size': 3, 'stride': 1, 'padding': 1}
|
|
conv = spectral_norm(nn.Conv2d(input_nc, input_nc, **kwargs), use_spect)
|
|
|
|
if type(norm_layer) == type(None):
|
|
self.model = nn.Sequential(conv, nonlinearity)
|
|
else:
|
|
self.model = nn.Sequential(conv, norm_layer(input_nc), nonlinearity)
|
|
|
|
def forward(self, x):
|
|
out = self.model(x)
|
|
return out
|
|
|
|
class FineADAINResBlock2d(nn.Module):
|
|
"""
|
|
Define an Residual block for different types
|
|
"""
|
|
def __init__(self, input_nc, feature_nc, norm_layer=nn.BatchNorm2d, nonlinearity=nn.LeakyReLU(), use_spect=False):
|
|
super(FineADAINResBlock2d, self).__init__()
|
|
|
|
kwargs = {'kernel_size': 3, 'stride': 1, 'padding': 1}
|
|
|
|
self.conv1 = spectral_norm(nn.Conv2d(input_nc, input_nc, **kwargs), use_spect)
|
|
self.conv2 = spectral_norm(nn.Conv2d(input_nc, input_nc, **kwargs), use_spect)
|
|
self.norm1 = ADAIN(input_nc, feature_nc)
|
|
self.norm2 = ADAIN(input_nc, feature_nc)
|
|
|
|
self.actvn = nonlinearity
|
|
|
|
|
|
def forward(self, x, z):
|
|
dx = self.actvn(self.norm1(self.conv1(x), z))
|
|
dx = self.norm2(self.conv2(x), z)
|
|
out = dx + x
|
|
return out
|
|
|
|
class FinalBlock2d(nn.Module):
|
|
"""
|
|
Define the output layer
|
|
"""
|
|
def __init__(self, input_nc, output_nc, use_spect=False, tanh_or_sigmoid='tanh'):
|
|
super(FinalBlock2d, self).__init__()
|
|
|
|
kwargs = {'kernel_size': 7, 'stride': 1, 'padding':3}
|
|
conv = spectral_norm(nn.Conv2d(input_nc, output_nc, **kwargs), use_spect)
|
|
|
|
if tanh_or_sigmoid == 'sigmoid':
|
|
out_nonlinearity = nn.Sigmoid()
|
|
else:
|
|
out_nonlinearity = nn.Tanh()
|
|
|
|
self.model = nn.Sequential(conv, out_nonlinearity)
|
|
def forward(self, x):
|
|
out = self.model(x)
|
|
return out |