567 lines
21 KiB
Python
567 lines
21 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from torch.nn.utils import spectral_norm
|
|
|
|
from .utils import init_weights, mean_variance_norm
|
|
|
|
|
|
class AdaIN(nn.Module):
|
|
def __init__(self):
|
|
super(AdaIN, self).__init__()
|
|
|
|
def forward(self, content, style, style_strength=1.0, eps=1e-5):
|
|
b, c, h, w = content.size()
|
|
|
|
content_std, content_mean = torch.std_mean(
|
|
content.view(b, c, -1), dim=2, keepdim=True
|
|
)
|
|
style_std, style_mean = torch.std_mean(
|
|
style.view(b, c, -1), dim=2, keepdim=True
|
|
)
|
|
|
|
normalized_content = (content.view(b, c, -1) - content_mean) / (
|
|
content_std + eps
|
|
)
|
|
|
|
stylized_content = (normalized_content * style_std) + style_mean
|
|
|
|
output = (
|
|
1 - style_strength
|
|
) * content + style_strength * stylized_content.view(b, c, h, w)
|
|
return output
|
|
|
|
|
|
class StylizedSANet(nn.Module):
|
|
def __init__(self, in_planes):
|
|
super(StylizedSANet, self).__init__()
|
|
self.f = nn.Conv2d(in_planes, in_planes, (1, 1))
|
|
self.g = nn.Conv2d(in_planes, in_planes, (1, 1))
|
|
self.h = nn.Conv2d(in_planes, in_planes, (1, 1))
|
|
self.sm = nn.Softmax(dim=-1)
|
|
self.out_conv = nn.Conv2d(in_planes, in_planes, (1, 1))
|
|
self.style_transform = AdaIN()
|
|
|
|
def forward(self, content, style):
|
|
F = self.f(mean_variance_norm(content))
|
|
G = self.g(mean_variance_norm(style))
|
|
H = self.h(style)
|
|
b, c, h, w = F.size()
|
|
F = F.view(b, -1, w * h).permute(0, 2, 1)
|
|
b, c, h, w = G.size()
|
|
G = G.view(b, -1, w * h)
|
|
S = torch.bmm(F, G)
|
|
S = self.sm(S)
|
|
b, c, h, w = H.size()
|
|
H = H.view(b, -1, w * h)
|
|
O = torch.bmm(H, S.permute(0, 2, 1))
|
|
b, c, h, w = content.size()
|
|
O = O.view(b, c, h, w)
|
|
O = self.out_conv(O)
|
|
# stylized_content = self.style_transform(content, style)
|
|
# O += stylized_content
|
|
O += content
|
|
return O
|
|
# 아니면 여기 adaptive하게 설정하는데 일단은 0.5 0.5 로 해볼까?
|
|
|
|
|
|
class Transform(nn.Module):
|
|
def __init__(self, in_planes):
|
|
super(Transform, self).__init__()
|
|
self.sanet4_1 = StylizedSANet(in_planes=in_planes)
|
|
self.sanet5_1 = StylizedSANet(in_planes=in_planes)
|
|
self.merge_conv_pad = nn.ReflectionPad2d((1, 1, 1, 1))
|
|
self.merge_conv = nn.Conv2d(in_planes, in_planes, (3, 3))
|
|
|
|
def forward(self, content4_1, style4_1, content5_1, style5_1):
|
|
# return self.merge_conv(self.merge_conv_pad(self.sanet4_1(content4_1, style4_1) + self.upsample5_1(self.sanet5_1(content5_1, style5_1))))
|
|
return self.merge_conv(
|
|
self.merge_conv_pad(
|
|
self.sanet4_1(content4_1, style4_1)
|
|
+ nn.functional.interpolate(
|
|
self.sanet5_1(content5_1, style5_1),
|
|
size=(
|
|
self.sanet4_1(content4_1, style4_1).size(2),
|
|
self.sanet4_1(content4_1, style4_1).size(3),
|
|
),
|
|
)
|
|
)
|
|
)
|
|
|
|
|
|
class AdaptiveMultiAdaAttN_v2(nn.Module):
|
|
def __init__(
|
|
self,
|
|
in_planes,
|
|
out_planes,
|
|
max_sample=256 * 256,
|
|
query_planes=None,
|
|
key_planes=None,
|
|
):
|
|
super(AdaptiveMultiAdaAttN_v2, self).__init__()
|
|
if key_planes is None:
|
|
key_planes = in_planes
|
|
self.f = nn.Conv2d(query_planes, key_planes, (1, 1))
|
|
self.g = nn.Conv2d(key_planes, key_planes, (1, 1))
|
|
self.h = nn.Conv2d(in_planes, out_planes, (1, 1))
|
|
self.sm = nn.Softmax(dim=-1)
|
|
self.out_conv = nn.Conv2d(in_planes, out_planes, (1, 1))
|
|
self.max_sample = max_sample
|
|
|
|
def forward(self, content, style, content_key, style_key, seed=None):
|
|
F = self.f(content_key)
|
|
G = self.g(style_key)
|
|
H = self.h(style)
|
|
b, _, h_g, w_g = G.size()
|
|
G = G.view(b, -1, w_g * h_g).contiguous()
|
|
b, _, h, w = F.size()
|
|
F = F.view(b, -1, w * h).permute(0, 2, 1)
|
|
S = torch.bmm(F, G)
|
|
# S: b, n_c, n_s
|
|
S = self.sm(S)
|
|
|
|
b, style_c, style_h, style_w = H.size()
|
|
H = torch.nn.functional.interpolate(H, (h_g, w_g), mode="bicubic")
|
|
if h_g * w_g > self.max_sample:
|
|
if seed is not None:
|
|
torch.manual_seed(seed)
|
|
index = torch.randperm(h_g * w_g).to(content.device)[: self.max_sample]
|
|
G = G[:, :, index]
|
|
style_flat = (
|
|
H.view(b, -1, h_g * w_g)[:, :, index].transpose(1, 2).contiguous()
|
|
)
|
|
else:
|
|
style_flat = H.view(b, -1, h_g * w_g).transpose(1, 2).contiguous()
|
|
|
|
# mean: b, n_c, c
|
|
mean = torch.bmm(S, style_flat)
|
|
# std: b, n_c, c
|
|
std = torch.sqrt(torch.relu(torch.bmm(S, style_flat**2) - mean**2))
|
|
# mean, std: b, c, h, w
|
|
_, _, ch, cw = content.size()
|
|
# mean = torch.nn.functional.interpolate(mean.view(b, style_h, style_w, style_c).permute(0, 3, 1, 2).contiguous(), (ch, cw))
|
|
# std = torch.nn.functional.interpolate(std.view(b, style_h, style_w, style_c).permute(0, 3, 1, 2).contiguous(), (ch, cw))
|
|
mean = mean.view(b, h, w, -1).permute(0, 3, 1, 2).contiguous()
|
|
std = std.view(b, h, w, -1).permute(0, 3, 1, 2).contiguous()
|
|
return std * mean_variance_norm(content) + mean, S
|
|
|
|
|
|
class AdaptiveMultiAttn_Transformer_v2(nn.Module):
|
|
def __init__(
|
|
self,
|
|
in_planes,
|
|
out_planes,
|
|
query_planes=None,
|
|
key_planes=None,
|
|
shallow_layer=False,
|
|
):
|
|
super(AdaptiveMultiAttn_Transformer_v2, self).__init__()
|
|
self.attn_adain_4_1 = AdaptiveMultiAdaAttN_v2(
|
|
in_planes=in_planes,
|
|
out_planes=out_planes,
|
|
query_planes=query_planes,
|
|
key_planes=key_planes,
|
|
)
|
|
self.attn_adain_5_1 = AdaptiveMultiAdaAttN_v2(
|
|
in_planes=in_planes,
|
|
out_planes=out_planes,
|
|
query_planes=query_planes,
|
|
key_planes=key_planes + 512,
|
|
)
|
|
# self.attn_adain_5_1 = AdaptiveMultiAdaAttN_v2(in_planes=in_planes, out_planes=out_planes, query_planes=query_planes+512, key_planes=key_planes+512)
|
|
self.upsample5_1 = nn.Upsample(scale_factor=2, mode="nearest")
|
|
self.merge_conv_pad = nn.ReflectionPad2d((1, 1, 1, 1))
|
|
self.merge_conv = nn.Conv2d(out_planes, out_planes, (3, 3))
|
|
|
|
def forward(
|
|
self,
|
|
content4_1,
|
|
style4_1,
|
|
content5_1,
|
|
style5_1,
|
|
content4_1_key,
|
|
style4_1_key,
|
|
content5_1_key,
|
|
style5_1_key,
|
|
seed=None,
|
|
):
|
|
feature_4_1, attn_4_1 = self.attn_adain_4_1(
|
|
content4_1, style4_1, content4_1_key, style4_1_key, seed=seed
|
|
)
|
|
feature_5_1, attn_5_1 = self.attn_adain_5_1(
|
|
content5_1, style5_1, content5_1_key, style5_1_key, seed=seed
|
|
)
|
|
# stylized_results = self.merge_conv(self.merge_conv_pad(feature_4_1 + self.upsample5_1(feature_5_1)))
|
|
|
|
stylized_results = self.merge_conv(
|
|
self.merge_conv_pad(
|
|
feature_4_1
|
|
+ nn.functional.interpolate(
|
|
feature_5_1, size=(feature_4_1.size(2), feature_4_1.size(3))
|
|
)
|
|
)
|
|
)
|
|
return stylized_results, feature_4_1, feature_5_1, attn_4_1, attn_5_1
|
|
|
|
|
|
class VGGEncoder(nn.Module):
|
|
def __init__(self, vgg_state_dict):
|
|
super(VGGEncoder, self).__init__()
|
|
|
|
self.pad = nn.ReflectionPad2d(1)
|
|
self.relu = nn.ReLU(inplace=True)
|
|
self.pool = nn.AvgPool2d(2)
|
|
self.maxpool1 = nn.MaxPool2d(kernel_size=2, stride=2, return_indices=False)
|
|
self.maxpool2 = nn.MaxPool2d(kernel_size=2, stride=2, return_indices=False)
|
|
self.maxpool3 = nn.MaxPool2d(kernel_size=2, stride=2, return_indices=False)
|
|
self.maxpool4 = nn.MaxPool2d(kernel_size=2, stride=2, return_indices=False)
|
|
|
|
###Level0###
|
|
self.conv0 = nn.Conv2d(3, 3, 1, 1, 0)
|
|
self.conv0.weight = nn.Parameter(vgg_state_dict["0.weight"])
|
|
self.conv0.bias = nn.Parameter(vgg_state_dict["0.bias"])
|
|
|
|
###Level1###
|
|
self.conv1_1 = nn.Conv2d(3, 64, 3, 1, 0)
|
|
self.conv1_1.weight = nn.Parameter(vgg_state_dict["2.weight"])
|
|
self.conv1_1.bias = nn.Parameter(vgg_state_dict["2.bias"])
|
|
|
|
self.conv1_2 = nn.Conv2d(64, 64, 3, 1, 0)
|
|
self.conv1_2.weight = nn.Parameter(vgg_state_dict["5.weight"])
|
|
self.conv1_2.bias = nn.Parameter(vgg_state_dict["5.bias"])
|
|
|
|
###Level2###
|
|
self.conv2_1 = nn.Conv2d(64, 128, 3, 1, 0)
|
|
self.conv2_1.weight = nn.Parameter(vgg_state_dict["9.weight"])
|
|
self.conv2_1.bias = nn.Parameter(vgg_state_dict["9.bias"])
|
|
|
|
self.conv2_2 = nn.Conv2d(128, 128, 3, 1, 0)
|
|
self.conv2_2.weight = nn.Parameter(vgg_state_dict["12.weight"])
|
|
self.conv2_2.bias = nn.Parameter(vgg_state_dict["12.bias"])
|
|
|
|
###Level3###
|
|
self.conv3_1 = nn.Conv2d(128, 256, 3, 1, 0)
|
|
self.conv3_1.weight = nn.Parameter(vgg_state_dict["16.weight"])
|
|
self.conv3_1.bias = nn.Parameter(vgg_state_dict["16.bias"])
|
|
|
|
self.conv3_2 = nn.Conv2d(256, 256, 3, 1, 0)
|
|
self.conv3_2.weight = nn.Parameter(vgg_state_dict["19.weight"])
|
|
self.conv3_2.bias = nn.Parameter(vgg_state_dict["19.bias"])
|
|
|
|
self.conv3_3 = nn.Conv2d(256, 256, 3, 1, 0)
|
|
self.conv3_3.weight = nn.Parameter(vgg_state_dict["22.weight"])
|
|
self.conv3_3.bias = nn.Parameter(vgg_state_dict["22.bias"])
|
|
|
|
self.conv3_4 = nn.Conv2d(256, 256, 3, 1, 0)
|
|
self.conv3_4.weight = nn.Parameter(vgg_state_dict["25.weight"])
|
|
self.conv3_4.bias = nn.Parameter(vgg_state_dict["25.bias"])
|
|
|
|
###Level4###
|
|
self.conv4_1 = nn.Conv2d(256, 512, 3, 1, 0)
|
|
self.conv4_1.weight = nn.Parameter(vgg_state_dict["29.weight"])
|
|
self.conv4_1.bias = nn.Parameter(vgg_state_dict["29.bias"])
|
|
|
|
self.conv4_2 = nn.Conv2d(512, 512, 3, 1, 0)
|
|
self.conv4_2.weight = nn.Parameter(vgg_state_dict["32.weight"])
|
|
self.conv4_2.bias = nn.Parameter(vgg_state_dict["32.bias"])
|
|
|
|
self.conv4_3 = nn.Conv2d(512, 512, 3, 1, 0)
|
|
self.conv4_3.weight = nn.Parameter(vgg_state_dict["35.weight"])
|
|
self.conv4_3.bias = nn.Parameter(vgg_state_dict["35.bias"])
|
|
|
|
self.conv4_4 = nn.Conv2d(512, 512, 3, 1, 0)
|
|
self.conv4_4.weight = nn.Parameter(vgg_state_dict["38.weight"])
|
|
self.conv4_4.bias = nn.Parameter(vgg_state_dict["38.bias"])
|
|
|
|
###Level5###
|
|
self.conv5_1 = nn.Conv2d(512, 512, 3, 1, 0)
|
|
self.conv5_1.weight = nn.Parameter(vgg_state_dict["42.weight"])
|
|
self.conv5_1.bias = nn.Parameter(vgg_state_dict["42.bias"])
|
|
|
|
def forward(self, x):
|
|
skips = {}
|
|
for level in [1, 2, 3, 4]:
|
|
x = self.encode(x, skips, level)
|
|
return x
|
|
|
|
def encode(self, x, skips):
|
|
is_maxpool = False
|
|
|
|
out = self.conv0(x)
|
|
out = self.relu(self.conv1_1(self.pad(out)))
|
|
skips["conv1_1"] = out
|
|
|
|
out = self.relu(self.conv1_2(self.pad(out)))
|
|
skips["conv1_2"] = out
|
|
resize_w, resize_h = out.size(2), out.size(3)
|
|
pooled_feature = self.pool(out)
|
|
|
|
# HH = out - F.interpolate(pooled_feature, scale_factor=2, mode='nearest')
|
|
HH = out - F.interpolate(
|
|
pooled_feature, size=[resize_w, resize_h], mode="nearest"
|
|
)
|
|
skips["pool1"] = HH
|
|
##################################
|
|
if is_maxpool:
|
|
pooled_feature = self.maxpool1(out)
|
|
|
|
out = self.relu(self.conv2_1(self.pad(pooled_feature)))
|
|
skips["conv2_1"] = out
|
|
|
|
out = self.relu(self.conv2_2(self.pad(out)))
|
|
skips["conv2_2"] = out
|
|
resize_w, resize_h = out.size(2), out.size(3)
|
|
pooled_feature = self.pool(out)
|
|
|
|
# HH = out - F.interpolate(pooled_feature, scale_factor=2, mode='nearest')
|
|
HH = out - F.interpolate(
|
|
pooled_feature, size=[resize_w, resize_h], mode="nearest"
|
|
)
|
|
skips["pool2"] = HH
|
|
##################################
|
|
if is_maxpool:
|
|
pooled_feature = self.maxpool2(out)
|
|
|
|
out = self.relu(self.conv3_1(self.pad(pooled_feature)))
|
|
skips["conv3_1"] = out
|
|
|
|
out = self.relu(self.conv3_2(self.pad(out)))
|
|
out = self.relu(self.conv3_3(self.pad(out)))
|
|
out = self.relu(self.conv3_4(self.pad(out)))
|
|
skips["conv3_4"] = out
|
|
resize_w, resize_h = out.size(2), out.size(3)
|
|
pooled_feature = self.pool(out)
|
|
# HH = out - F.interpolate(pooled_feature, scale_factor=2, mode='nearest')
|
|
HH = out - F.interpolate(
|
|
pooled_feature, size=[resize_w, resize_h], mode="nearest"
|
|
)
|
|
skips["pool3"] = HH
|
|
|
|
##################################
|
|
if is_maxpool:
|
|
pooled_feature = self.maxpool3(out)
|
|
|
|
out = self.relu(self.conv4_1(self.pad(pooled_feature)))
|
|
skips["conv4_1"] = out
|
|
|
|
out = self.relu(self.conv4_2(self.pad(out)))
|
|
out = self.relu(self.conv4_3(self.pad(out)))
|
|
out = self.relu(self.conv4_4(self.pad(out)))
|
|
skips["conv4_4"] = out
|
|
resize_w, resize_h = out.size(2), out.size(3)
|
|
pooled_feature = self.pool(out)
|
|
HH = out - F.interpolate(
|
|
pooled_feature, size=[resize_w, resize_h], mode="nearest"
|
|
)
|
|
skips["pool4"] = HH
|
|
|
|
####################################
|
|
if is_maxpool:
|
|
pooled_feature = self.maxpool4(out)
|
|
out = self.relu(self.conv5_1(self.pad(pooled_feature)))
|
|
skips["conv5_1"] = out
|
|
|
|
return out
|
|
|
|
def get_features(self, x, level):
|
|
is_maxpool = False
|
|
|
|
out = self.conv0(x)
|
|
out = self.relu(self.conv1_1(self.pad(out)))
|
|
if level == 1:
|
|
return out
|
|
|
|
out = self.relu(self.conv1_2(self.pad(out)))
|
|
pooled_feature = self.pool(out)
|
|
##################################
|
|
if is_maxpool:
|
|
pooled_feature = self.maxpool1(out)
|
|
out = self.relu(self.conv2_1(self.pad(pooled_feature)))
|
|
if level == 2:
|
|
return out
|
|
|
|
out = self.relu(self.conv2_2(self.pad(out)))
|
|
pooled_feature = self.pool(out)
|
|
##################################
|
|
if is_maxpool:
|
|
pooled_feature = self.maxpool2(out)
|
|
out = self.relu(self.conv3_1(self.pad(pooled_feature)))
|
|
if level == 3:
|
|
return out
|
|
|
|
out = self.relu(self.conv3_2(self.pad(out)))
|
|
out = self.relu(self.conv3_3(self.pad(out)))
|
|
out = self.relu(self.conv3_4(self.pad(out)))
|
|
pooled_feature = self.pool(out)
|
|
##################################
|
|
if is_maxpool:
|
|
pooled_feature = self.maxpool3(out)
|
|
out = self.relu(self.conv4_1(self.pad(pooled_feature)))
|
|
if level == 4:
|
|
return out
|
|
|
|
out = self.relu(self.conv4_2(self.pad(out)))
|
|
out = self.relu(self.conv4_3(self.pad(out)))
|
|
out = self.relu(self.conv4_4(self.pad(out)))
|
|
pooled_feature = self.pool(out)
|
|
####################################
|
|
if is_maxpool:
|
|
pooled_feature = self.maxpool4(out)
|
|
out = self.relu(self.conv5_1(self.pad(pooled_feature)))
|
|
if level == 5:
|
|
return out
|
|
|
|
|
|
class VGGDecoder(nn.Module):
|
|
def __init__(self):
|
|
super(VGGDecoder, self).__init__()
|
|
|
|
self.pad = nn.ReflectionPad2d(1)
|
|
self.relu = nn.ReLU(inplace=False)
|
|
self.adain = AdaIN()
|
|
|
|
# self.conv5_1 = nn.Conv2d(512, 512, 3, 1, 0)
|
|
# self.conv4_4 = nn.Conv2d(512, 512, 3, 1, 0)
|
|
# self.conv4_3 = nn.Conv2d(512, 512, 3, 1, 0)
|
|
# self.conv4_2 = nn.Conv2d(512, 512, 3, 1, 0)
|
|
self.conv4_1 = nn.Conv2d(512, 256, 3, 1, 0)
|
|
self.conv3_4 = nn.Conv2d(256, 256, 3, 1, 0)
|
|
self.conv3_3 = nn.Conv2d(256, 256, 3, 1, 0)
|
|
self.conv3_2 = nn.Conv2d(256, 256, 3, 1, 0)
|
|
self.conv3_1 = nn.Conv2d(256, 128, 3, 1, 0)
|
|
self.conv2_2 = nn.Conv2d(128, 128, 3, 1, 0)
|
|
self.conv2_1 = nn.Conv2d(128, 64, 3, 1, 0)
|
|
self.conv1_2 = nn.Conv2d(64, 64, 3, 1, 0)
|
|
self.conv1_1 = nn.Conv2d(64, 3, 3, 1, 0)
|
|
|
|
init_weights(self)
|
|
|
|
def forward(self, x, skips):
|
|
x = self.decode(x, skips)
|
|
|
|
def decode(self, stylized_feat, content_skips, style_skips):
|
|
out = self.relu(self.conv4_1(self.pad(stylized_feat)))
|
|
resize_w, resize_h = (
|
|
content_skips["conv3_4"].size(2),
|
|
content_skips["conv3_4"].size(3),
|
|
)
|
|
unpooled_feat = F.interpolate(out, size=[resize_w, resize_h], mode="nearest")
|
|
out = unpooled_feat
|
|
|
|
##여기에 skip_conneciton 넣어서 stylization boost 하기 (local pattern)
|
|
out = self.relu(self.conv3_4(self.pad(out)))
|
|
out = self.relu(self.conv3_3(self.pad(out)))
|
|
out = self.relu(self.conv3_2(self.pad(out)))
|
|
# out = feature_wct_simple(out, style_skips['conv3_1'])
|
|
|
|
out = self.relu(self.conv3_1(self.pad(out)))
|
|
resize_w, resize_h = (
|
|
content_skips["conv2_2"].size(2),
|
|
content_skips["conv2_2"].size(3),
|
|
)
|
|
unpooled_feat = F.interpolate(out, size=[resize_w, resize_h], mode="nearest")
|
|
out = unpooled_feat
|
|
|
|
out = self.relu(self.conv2_2(self.pad(out)))
|
|
# out = feature_wct_simple(out, style_skips['conv2_1'])
|
|
|
|
out = self.relu(self.conv2_1(self.pad(out)))
|
|
resize_w, resize_h = (
|
|
content_skips["conv1_2"].size(2),
|
|
content_skips["conv1_2"].size(3),
|
|
)
|
|
unpooled_feat = F.interpolate(out, size=[resize_w, resize_h], mode="nearest")
|
|
out = unpooled_feat
|
|
|
|
out = self.relu(self.conv1_2(self.pad(out)))
|
|
# out = feature_wct_simple(out, style_skips['conv1_1'])
|
|
out = self.conv1_1(self.pad(out))
|
|
return out
|
|
|
|
def reconstruct(self, x):
|
|
out = self.relu(self.conv4_1(self.pad(x)))
|
|
out = F.interpolate(out, size=[out.size(2), out.size(3)], mode="nearest")
|
|
out = self.relu(self.conv3_4(self.pad(out)))
|
|
out = self.relu(self.conv3_3(self.pad(out)))
|
|
out = self.relu(self.conv3_2(self.pad(out)))
|
|
out = self.relu(self.conv3_1(self.pad(out)))
|
|
out = F.interpolate(out, size=[out.size(2), out.size(3)], mode="nearest")
|
|
out = self.relu(self.conv2_2(self.pad(out)))
|
|
out = self.relu(self.conv2_1(self.pad(out)))
|
|
out = F.interpolate(out, size=[out.size(2), out.size(3)], mode="nearest")
|
|
out = self.relu(self.conv1_2(self.pad(out)))
|
|
out = self.conv1_1(self.pad(out))
|
|
return out
|
|
|
|
|
|
#################################
|
|
#####MultiScaleDiscriminator#####
|
|
#################################
|
|
class MultiScaleImageDiscriminator(nn.Module):
|
|
def __init__(self, nc=3, ndf=64):
|
|
super(MultiScaleImageDiscriminator, self).__init__()
|
|
self.nc = nc
|
|
self.output_dim = 1
|
|
|
|
self.conv1 = nn.Sequential(
|
|
# 256->128
|
|
spectral_norm(nn.Conv2d(self.nc, ndf, 4, 2, 1, bias=False)),
|
|
nn.LeakyReLU(0.2, inplace=True),
|
|
)
|
|
|
|
self.skip_out1 = nn.Sequential(
|
|
nn.Conv2d(ndf, 32, 3, 1, 1, bias=False),
|
|
nn.LeakyReLU(0.2, inplace=True),
|
|
nn.AdaptiveAvgPool2d(8),
|
|
nn.Conv2d(32, self.output_dim, 3, 1, 1, bias=False),
|
|
)
|
|
|
|
self.conv2 = nn.Sequential(
|
|
# 128->64
|
|
spectral_norm(nn.Conv2d(ndf, ndf * 2, 4, 2, 1, bias=False)),
|
|
# nn.BatchNorm2d(ndf*2),
|
|
nn.InstanceNorm2d(ndf * 2),
|
|
nn.LeakyReLU(0.2, inplace=True),
|
|
)
|
|
|
|
self.conv3 = nn.Sequential(
|
|
# 64->32
|
|
spectral_norm(nn.Conv2d(ndf * 2, ndf * 4, 4, 2, 1, bias=False)),
|
|
# nn.BatchNorm2d(ndf*4),
|
|
nn.InstanceNorm2d(ndf * 4),
|
|
nn.LeakyReLU(0.2, inplace=True),
|
|
)
|
|
|
|
self.skip_out2 = nn.Sequential(
|
|
nn.Conv2d(ndf * 4, 32, 3, 1, 1, bias=False),
|
|
nn.LeakyReLU(0.2, inplace=True),
|
|
nn.AdaptiveAvgPool2d(8),
|
|
nn.Conv2d(32, self.output_dim, 3, 1, 1, bias=False),
|
|
)
|
|
|
|
self.conv4 = nn.Sequential(
|
|
# 32->16
|
|
spectral_norm(nn.Conv2d(ndf * 4, ndf * 8, 4, 2, 1, bias=False)),
|
|
# nn.BatchNorm2d(ndf*8),
|
|
nn.InstanceNorm2d(ndf * 8),
|
|
nn.LeakyReLU(0.2, inplace=True),
|
|
)
|
|
|
|
self.conv5 = nn.Sequential(
|
|
# 16->8
|
|
spectral_norm(nn.Conv2d(ndf * 8, self.output_dim, 4, 2, 1, bias=False)),
|
|
)
|
|
|
|
init_weights(self)
|
|
|
|
def forward(self, input):
|
|
out = self.conv1(input)
|
|
skip_out1 = self.skip_out1(out)
|
|
out = self.conv3(self.conv2(out))
|
|
skip_out2 = self.skip_out2(out)
|
|
out = self.conv5(self.conv4(out))
|
|
out = ((out + skip_out1 + skip_out2) * 1 / 3).squeeze()
|
|
# out = ((out+skip_out1+skip_out2)).squeeze()
|
|
|
|
return out
|