|
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__() |
|
|
|
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} |
|
|
|
|
|
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)) |
|
|
|
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): |
|
|
|
|
|
normalized = self.param_free_norm(x) |
|
|
|
|
|
feature = feature.view(feature.size(0), -1) |
|
actv = self.mlp_shared(feature) |
|
gamma = self.mlp_gamma(actv) |
|
beta = self.mlp_beta(actv) |
|
|
|
|
|
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 |
