update training code
This commit is contained in:
zhangmozhe
@@ -0,0 +1,303 @@
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import torchvision
|
||||
from utils.util import feature_normalize
|
||||
|
||||
postpa = torchvision.transforms.Compose(
|
||||
[
|
||||
torchvision.transforms.Lambda(lambda x: x.mul_(1.0 / 255)),
|
||||
torchvision.transforms.Normalize(
|
||||
mean=[-0.40760392, -0.45795686, -0.48501961], std=[1, 1, 1] # add imagenet mean
|
||||
),
|
||||
torchvision.transforms.Lambda(lambda x: x[torch.LongTensor([2, 1, 0])]), # turn to RGB
|
||||
]
|
||||
)
|
||||
postpb = torchvision.transforms.Compose([torchvision.transforms.ToPILImage()])
|
||||
|
||||
|
||||
def post_processing(tensor):
|
||||
t = postpa(tensor) # denormalize the image since the optimized tensor is the normalized one
|
||||
t[t > 1] = 1
|
||||
t[t < 0] = 0
|
||||
img = postpb(t)
|
||||
img = np.array(img)
|
||||
return img
|
||||
|
||||
|
||||
class ContextualLoss(nn.Module):
|
||||
"""
|
||||
input is Al, Bl, channel = 1, range ~ [0, 255]
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super(ContextualLoss, self).__init__()
|
||||
return None
|
||||
|
||||
def forward(self, X_features, Y_features, h=0.1, feature_centering=True):
|
||||
"""
|
||||
X_features&Y_features are are feature vectors or feature 2d array
|
||||
h: bandwidth
|
||||
return the per-sample loss
|
||||
"""
|
||||
batch_size = X_features.shape[0]
|
||||
feature_depth = X_features.shape[1]
|
||||
feature_size = X_features.shape[2]
|
||||
|
||||
# to normalized feature vectors
|
||||
if feature_centering:
|
||||
X_features = X_features - Y_features.view(batch_size, feature_depth, -1).mean(dim=-1).unsqueeze(
|
||||
dim=-1
|
||||
).unsqueeze(dim=-1)
|
||||
Y_features = Y_features - Y_features.view(batch_size, feature_depth, -1).mean(dim=-1).unsqueeze(
|
||||
dim=-1
|
||||
).unsqueeze(dim=-1)
|
||||
X_features = feature_normalize(X_features).view(
|
||||
batch_size, feature_depth, -1
|
||||
) # batch_size * feature_depth * feature_size^2
|
||||
Y_features = feature_normalize(Y_features).view(
|
||||
batch_size, feature_depth, -1
|
||||
) # batch_size * feature_depth * feature_size^2
|
||||
|
||||
# conine distance = 1 - similarity
|
||||
X_features_permute = X_features.permute(0, 2, 1) # batch_size * feature_size^2 * feature_depth
|
||||
d = 1 - torch.matmul(X_features_permute, Y_features) # batch_size * feature_size^2 * feature_size^2
|
||||
|
||||
# normalized distance: dij_bar
|
||||
d_norm = d / (torch.min(d, dim=-1, keepdim=True)[0] + 1e-5) # batch_size * feature_size^2 * feature_size^2
|
||||
|
||||
# pairwise affinity
|
||||
w = torch.exp((1 - d_norm) / h)
|
||||
A_ij = w / torch.sum(w, dim=-1, keepdim=True)
|
||||
|
||||
# contextual loss per sample
|
||||
CX = torch.mean(torch.max(A_ij, dim=1)[0], dim=-1)
|
||||
return -torch.log(CX)
|
||||
|
||||
|
||||
class ContextualLoss_forward(nn.Module):
|
||||
"""
|
||||
input is Al, Bl, channel = 1, range ~ [0, 255]
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super(ContextualLoss_forward, self).__init__()
|
||||
return None
|
||||
|
||||
def forward(self, X_features, Y_features, h=0.1, feature_centering=True):
|
||||
"""
|
||||
X_features&Y_features are are feature vectors or feature 2d array
|
||||
h: bandwidth
|
||||
return the per-sample loss
|
||||
"""
|
||||
batch_size = X_features.shape[0]
|
||||
feature_depth = X_features.shape[1]
|
||||
feature_size = X_features.shape[2]
|
||||
|
||||
# to normalized feature vectors
|
||||
if feature_centering:
|
||||
X_features = X_features - Y_features.view(batch_size, feature_depth, -1).mean(dim=-1).unsqueeze(
|
||||
dim=-1
|
||||
).unsqueeze(dim=-1)
|
||||
Y_features = Y_features - Y_features.view(batch_size, feature_depth, -1).mean(dim=-1).unsqueeze(
|
||||
dim=-1
|
||||
).unsqueeze(dim=-1)
|
||||
X_features = feature_normalize(X_features).view(
|
||||
batch_size, feature_depth, -1
|
||||
) # batch_size * feature_depth * feature_size^2
|
||||
Y_features = feature_normalize(Y_features).view(
|
||||
batch_size, feature_depth, -1
|
||||
) # batch_size * feature_depth * feature_size^2
|
||||
|
||||
# conine distance = 1 - similarity
|
||||
X_features_permute = X_features.permute(0, 2, 1) # batch_size * feature_size^2 * feature_depth
|
||||
d = 1 - torch.matmul(X_features_permute, Y_features) # batch_size * feature_size^2 * feature_size^2
|
||||
|
||||
# normalized distance: dij_bar
|
||||
d_norm = d / (torch.min(d, dim=-1, keepdim=True)[0] + 1e-5) # batch_size * feature_size^2 * feature_size^2
|
||||
|
||||
# pairwise affinity
|
||||
w = torch.exp((1 - d_norm) / h)
|
||||
A_ij = w / torch.sum(w, dim=-1, keepdim=True)
|
||||
|
||||
# contextual loss per sample
|
||||
CX = torch.mean(torch.max(A_ij, dim=-1)[0], dim=1)
|
||||
return -torch.log(CX)
|
||||
|
||||
|
||||
class ContextualLoss_complex(nn.Module):
|
||||
"""
|
||||
input is Al, Bl, channel = 1, range ~ [0, 255]
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super(ContextualLoss_complex, self).__init__()
|
||||
return None
|
||||
|
||||
def forward(self, X_features, Y_features, h=0.1, patch_size=1, direction="forward"):
|
||||
"""
|
||||
X_features&Y_features are are feature vectors or feature 2d array
|
||||
h: bandwidth
|
||||
return the per-sample loss
|
||||
"""
|
||||
batch_size = X_features.shape[0]
|
||||
feature_depth = X_features.shape[1]
|
||||
feature_size = X_features.shape[2]
|
||||
|
||||
# to normalized feature vectors
|
||||
X_features = X_features - Y_features.view(batch_size, feature_depth, -1).mean(dim=-1).unsqueeze(
|
||||
dim=-1
|
||||
).unsqueeze(dim=-1)
|
||||
Y_features = Y_features - Y_features.view(batch_size, feature_depth, -1).mean(dim=-1).unsqueeze(
|
||||
dim=-1
|
||||
).unsqueeze(dim=-1)
|
||||
X_features = feature_normalize(X_features) # batch_size * feature_depth * feature_size^2
|
||||
Y_features = feature_normalize(Y_features) # batch_size * feature_depth * feature_size^2
|
||||
|
||||
# to normalized feature vectors
|
||||
X_features = F.unfold(
|
||||
X_features, kernel_size=(patch_size, patch_size), stride=(1, 1), padding=(patch_size // 2, patch_size // 2)
|
||||
) # batch_size * feature_depth_new * feature_size^2
|
||||
Y_features = F.unfold(
|
||||
Y_features, kernel_size=(patch_size, patch_size), stride=(1, 1), padding=(patch_size // 2, patch_size // 2)
|
||||
) # batch_size * feature_depth_new * feature_size^2
|
||||
|
||||
# conine distance = 1 - similarity
|
||||
X_features_permute = X_features.permute(0, 2, 1) # batch_size * feature_size^2 * feature_depth
|
||||
d = 1 - torch.matmul(X_features_permute, Y_features) # batch_size * feature_size^2 * feature_size^2
|
||||
|
||||
# normalized distance: dij_bar
|
||||
d_norm = d / (torch.min(d, dim=-1, keepdim=True)[0] + 1e-5) # batch_size * feature_size^2 * feature_size^2
|
||||
|
||||
# pairwise affinity
|
||||
w = torch.exp((1 - d_norm) / h)
|
||||
A_ij = w / torch.sum(w, dim=-1, keepdim=True)
|
||||
|
||||
# contextual loss per sample
|
||||
if direction == "forward":
|
||||
CX = torch.mean(torch.max(A_ij, dim=-1)[0], dim=1)
|
||||
else:
|
||||
CX = torch.mean(torch.max(A_ij, dim=1)[0], dim=-1)
|
||||
|
||||
return -torch.log(CX)
|
||||
|
||||
|
||||
class ChamferDistance_patch_loss(nn.Module):
|
||||
"""
|
||||
input is Al, Bl, channel = 1, range ~ [0, 255]
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super(ChamferDistance_patch_loss, self).__init__()
|
||||
return None
|
||||
|
||||
def forward(self, X_features, Y_features, patch_size=3, image_x=None, image_y=None, h=0.1, Y_features_in=None):
|
||||
"""
|
||||
X_features&Y_features are are feature vectors or feature 2d array
|
||||
h: bandwidth
|
||||
return the per-sample loss
|
||||
"""
|
||||
batch_size = X_features.shape[0]
|
||||
feature_depth = X_features.shape[1]
|
||||
feature_size = X_features.shape[2]
|
||||
|
||||
# to normalized feature vectors
|
||||
X_features = F.unfold(
|
||||
X_features, kernel_size=(patch_size, patch_size), stride=(1, 1), padding=(patch_size // 2, patch_size // 2)
|
||||
) # batch_size, feature_depth_new * feature_size^2
|
||||
Y_features = F.unfold(
|
||||
Y_features, kernel_size=(patch_size, patch_size), stride=(1, 1), padding=(patch_size // 2, patch_size // 2)
|
||||
) # batch_size, feature_depth_new * feature_size^2
|
||||
|
||||
if image_x is not None and image_y is not None:
|
||||
image_x = torch.nn.functional.interpolate(image_x, size=(feature_size, feature_size), mode="bilinear").view(
|
||||
batch_size, 3, -1
|
||||
)
|
||||
image_y = torch.nn.functional.interpolate(image_y, size=(feature_size, feature_size), mode="bilinear").view(
|
||||
batch_size, 3, -1
|
||||
)
|
||||
|
||||
X_features_permute = X_features.permute(0, 2, 1) # batch_size * feature_size^2 * feature_depth
|
||||
similarity_matrix = torch.matmul(X_features_permute, Y_features) # batch_size * feature_size^2 * feature_size^2
|
||||
NN_index = similarity_matrix.max(dim=-1, keepdim=True)[1].squeeze()
|
||||
|
||||
if Y_features_in is not None:
|
||||
loss = torch.mean((X_features - Y_features_in.detach()) ** 2)
|
||||
Y_features_in = Y_features_in.detach()
|
||||
else:
|
||||
loss = torch.mean((X_features - Y_features[:, :, NN_index].detach()) ** 2)
|
||||
Y_features_in = Y_features[:, :, NN_index].detach()
|
||||
|
||||
# re-arrange image
|
||||
if image_x is not None and image_y is not None:
|
||||
image_y_rearrange = image_y[:, :, NN_index]
|
||||
image_y_rearrange = image_y_rearrange.view(batch_size, 3, feature_size, feature_size)
|
||||
image_x = image_x.view(batch_size, 3, feature_size, feature_size)
|
||||
image_y = image_y.view(batch_size, 3, feature_size, feature_size)
|
||||
|
||||
return loss
|
||||
|
||||
|
||||
class ChamferDistance_loss(nn.Module):
|
||||
"""
|
||||
input is Al, Bl, channel = 1, range ~ [0, 255]
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super(ChamferDistance_loss, self).__init__()
|
||||
return None
|
||||
|
||||
def forward(self, X_features, Y_features, image_x, image_y, h=0.1, Y_features_in=None):
|
||||
"""
|
||||
X_features&Y_features are are feature vectors or feature 2d array
|
||||
h: bandwidth
|
||||
return the per-sample loss
|
||||
"""
|
||||
batch_size = X_features.shape[0]
|
||||
feature_depth = X_features.shape[1]
|
||||
feature_size = X_features.shape[2]
|
||||
|
||||
# to normalized feature vectors
|
||||
X_features = feature_normalize(X_features).view(
|
||||
batch_size, feature_depth, -1
|
||||
) # batch_size * feature_depth * feature_size^2
|
||||
Y_features = feature_normalize(Y_features).view(
|
||||
batch_size, feature_depth, -1
|
||||
) # batch_size * feature_depth * feature_size^2
|
||||
image_x = torch.nn.functional.interpolate(image_x, size=(feature_size, feature_size), mode="bilinear").view(
|
||||
batch_size, 3, -1
|
||||
)
|
||||
image_y = torch.nn.functional.interpolate(image_y, size=(feature_size, feature_size), mode="bilinear").view(
|
||||
batch_size, 3, -1
|
||||
)
|
||||
|
||||
X_features_permute = X_features.permute(0, 2, 1) # batch_size * feature_size^2 * feature_depth
|
||||
similarity_matrix = torch.matmul(X_features_permute, Y_features) # batch_size * feature_size^2 * feature_size^2
|
||||
NN_index = similarity_matrix.max(dim=-1, keepdim=True)[1].squeeze()
|
||||
if Y_features_in is not None:
|
||||
loss = torch.mean((X_features - Y_features_in.detach()) ** 2)
|
||||
Y_features_in = Y_features_in.detach()
|
||||
else:
|
||||
loss = torch.mean((X_features - Y_features[:, :, NN_index].detach()) ** 2)
|
||||
Y_features_in = Y_features[:, :, NN_index].detach()
|
||||
|
||||
# re-arrange image
|
||||
image_y_rearrange = image_y[:, :, NN_index]
|
||||
image_y_rearrange = image_y_rearrange.view(batch_size, 3, feature_size, feature_size)
|
||||
image_x = image_x.view(batch_size, 3, feature_size, feature_size)
|
||||
image_y = image_y.view(batch_size, 3, feature_size, feature_size)
|
||||
|
||||
return loss, Y_features_in, X_features
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
contextual_loss = ContextualLoss()
|
||||
batch_size = 32
|
||||
feature_depth = 8
|
||||
feature_size = 16
|
||||
X_features = torch.zeros(batch_size, feature_depth, feature_size, feature_size)
|
||||
Y_features = torch.zeros(batch_size, feature_depth, feature_size, feature_size)
|
||||
|
||||
cx_loss = contextual_loss(X_features, Y_features, 1)
|
||||
print(cx_loss)
|
||||
+207
-207
@@ -101,61 +101,13 @@ class Self_Attention(nn.Module):
|
||||
return out
|
||||
|
||||
|
||||
class Discriminator_WITHOUT_FC_x64_video(nn.Module):
|
||||
class Discriminator_x64(nn.Module):
|
||||
"""
|
||||
Discriminative Network
|
||||
"""
|
||||
|
||||
def __init__(self, in_size=6, ndf=64):
|
||||
super(Discriminator_WITHOUT_FC_x64_video, self).__init__()
|
||||
self.in_size = in_size
|
||||
self.ndf = ndf
|
||||
|
||||
self.layer1 = nn.Sequential(
|
||||
# input size is in_size x 64 x 64
|
||||
SpectralNorm(nn.Conv2d(self.in_size, self.ndf, 4, 2, 1)),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
)
|
||||
self.layer2 = nn.Sequential(
|
||||
# state size: ndf x 32 x 32
|
||||
SpectralNorm(nn.Conv2d(self.ndf, self.ndf * 2, 4, 2, 1)),
|
||||
nn.InstanceNorm2d(self.ndf * 2),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
)
|
||||
self.layer3 = nn.Sequential(
|
||||
# state size: (ndf * 2) x 16 x 16
|
||||
SpectralNorm(nn.Conv2d(self.ndf * 2, self.ndf * 4, 4, 2, 1)),
|
||||
nn.InstanceNorm2d(self.ndf * 4),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
)
|
||||
self.layer4 = nn.Sequential(
|
||||
# state size: (ndf * 4) x 8 x 8
|
||||
# Self_Attention(self.ndf * 4, 'relu'),
|
||||
SpectralNorm(nn.Conv2d(self.ndf * 4, self.ndf * 8, 4, 2, 1)),
|
||||
nn.InstanceNorm2d(self.ndf * 8),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
)
|
||||
|
||||
self.last = SpectralNorm(nn.Conv2d(self.ndf * 8, 1, [3, 6], 1, 0))
|
||||
|
||||
def forward(self, input):
|
||||
feature1 = self.layer1(input)
|
||||
feature2 = self.layer2(feature1)
|
||||
feature3 = self.layer3(feature2)
|
||||
feature4 = self.layer4(feature3)
|
||||
output = self.last(feature4)
|
||||
output = F.avg_pool2d(output, output.size()[2:]).view(output.size()[0], -1)
|
||||
|
||||
return output, feature4
|
||||
|
||||
|
||||
class Discriminator_WITHOUT_FC_x64_video_BIG(nn.Module):
|
||||
"""
|
||||
Discriminative Network
|
||||
"""
|
||||
|
||||
def __init__(self, in_size=6, ndf=64):
|
||||
super(Discriminator_WITHOUT_FC_x64_video_BIG, self).__init__()
|
||||
super(Discriminator_x64, self).__init__()
|
||||
self.in_size = in_size
|
||||
self.ndf = ndf
|
||||
|
||||
@@ -205,189 +157,237 @@ class Discriminator_WITHOUT_FC_x64_video_BIG(nn.Module):
|
||||
return output, feature4
|
||||
|
||||
|
||||
class Discriminator_WITHOUT_FC_x64(nn.Module):
|
||||
"""
|
||||
Discriminative Network
|
||||
"""
|
||||
# class Discriminator_WITHOUT_FC_x64_video(nn.Module):
|
||||
# """
|
||||
# Discriminative Network
|
||||
# """
|
||||
|
||||
def __init__(self, in_size=3, ndf=64):
|
||||
super(Discriminator_WITHOUT_FC_x64, self).__init__()
|
||||
self.in_size = in_size
|
||||
self.ndf = ndf
|
||||
# def __init__(self, in_size=6, ndf=64):
|
||||
# super(Discriminator_WITHOUT_FC_x64_video, self).__init__()
|
||||
# self.in_size = in_size
|
||||
# self.ndf = ndf
|
||||
|
||||
self.main = nn.Sequential(
|
||||
# input size is in_size x 64 x 64
|
||||
SpectralNorm(nn.Conv2d(self.in_size, self.ndf, 4, 2, 1)),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
# state size: ndf x 32 x 32
|
||||
Self_Attention(self.ndf),
|
||||
SpectralNorm(nn.Conv2d(self.ndf, self.ndf * 2, 4, 2, 1)),
|
||||
nn.InstanceNorm2d(self.ndf * 2),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
# state size: (ndf * 2) x 16 x 16
|
||||
SpectralNorm(nn.Conv2d(self.ndf * 2, self.ndf * 4, 4, 2, 1)),
|
||||
nn.InstanceNorm2d(self.ndf * 4),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
# state size: (ndf * 4) x 8 x 8
|
||||
SpectralNorm(nn.Conv2d(self.ndf * 4, self.ndf * 8, 4, 2, 1)),
|
||||
nn.InstanceNorm2d(self.ndf * 8),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
# sate size: (ndf * 8) x 4 x 4
|
||||
)
|
||||
# self.layer1 = nn.Sequential(
|
||||
# # input size is in_size x 64 x 64
|
||||
# SpectralNorm(nn.Conv2d(self.in_size, self.ndf, 4, 2, 1)),
|
||||
# nn.LeakyReLU(0.2, inplace=True),
|
||||
# )
|
||||
# self.layer2 = nn.Sequential(
|
||||
# # state size: ndf x 32 x 32
|
||||
# SpectralNorm(nn.Conv2d(self.ndf, self.ndf * 2, 4, 2, 1)),
|
||||
# nn.InstanceNorm2d(self.ndf * 2),
|
||||
# nn.LeakyReLU(0.2, inplace=True),
|
||||
# )
|
||||
# self.layer3 = nn.Sequential(
|
||||
# # state size: (ndf * 2) x 16 x 16
|
||||
# SpectralNorm(nn.Conv2d(self.ndf * 2, self.ndf * 4, 4, 2, 1)),
|
||||
# nn.InstanceNorm2d(self.ndf * 4),
|
||||
# nn.LeakyReLU(0.2, inplace=True),
|
||||
# )
|
||||
# self.layer4 = nn.Sequential(
|
||||
# # state size: (ndf * 4) x 8 x 8
|
||||
# # Self_Attention(self.ndf * 4, 'relu'),
|
||||
# SpectralNorm(nn.Conv2d(self.ndf * 4, self.ndf * 8, 4, 2, 1)),
|
||||
# nn.InstanceNorm2d(self.ndf * 8),
|
||||
# nn.LeakyReLU(0.2, inplace=True),
|
||||
# )
|
||||
|
||||
self.last = SpectralNorm(nn.Conv2d(self.ndf * 8, 1, 4, 1, 0))
|
||||
# self.last = SpectralNorm(nn.Conv2d(self.ndf * 8, 1, [3, 6], 1, 0))
|
||||
|
||||
def forward(self, input):
|
||||
# def forward(self, input):
|
||||
# feature1 = self.layer1(input)
|
||||
# feature2 = self.layer2(feature1)
|
||||
# feature3 = self.layer3(feature2)
|
||||
# feature4 = self.layer4(feature3)
|
||||
# output = self.last(feature4)
|
||||
# output = F.avg_pool2d(output, output.size()[2:]).view(output.size()[0], -1)
|
||||
|
||||
feature = self.main(input)
|
||||
output = self.last(feature)
|
||||
output = F.avg_pool2d(output, output.size()[2:]).view(output.size()[0], -1)
|
||||
|
||||
return output, feature
|
||||
# return output, feature4
|
||||
|
||||
|
||||
class Discriminator_WITHOUT_FC_x64_BIG(nn.Module):
|
||||
"""
|
||||
Discriminative Network
|
||||
"""
|
||||
# class Discriminator_WITHOUT_FC_x64(nn.Module):
|
||||
# """
|
||||
# Discriminative Network
|
||||
# """
|
||||
|
||||
def __init__(self, in_size=3, ndf=64):
|
||||
super(Discriminator_WITHOUT_FC_x64_BIG, self).__init__()
|
||||
self.in_size = in_size
|
||||
self.ndf = ndf
|
||||
# def __init__(self, in_size=3, ndf=64):
|
||||
# super(Discriminator_WITHOUT_FC_x64, self).__init__()
|
||||
# self.in_size = in_size
|
||||
# self.ndf = ndf
|
||||
|
||||
self.main = nn.Sequential(
|
||||
# 256 X 256
|
||||
SpectralNorm(nn.Conv2d(self.in_size, self.ndf, 4, 2, 1)),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
# 128 X 128
|
||||
SpectralNorm(nn.Conv2d(self.ndf, self.ndf, 4, 2, 1)),
|
||||
nn.InstanceNorm2d(self.ndf),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
# 64 X 64
|
||||
Self_Attention(self.ndf),
|
||||
SpectralNorm(nn.Conv2d(self.ndf, self.ndf * 2, 4, 2, 1)),
|
||||
nn.InstanceNorm2d(self.ndf * 2),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
# 32 X 32
|
||||
SpectralNorm(nn.Conv2d(self.ndf * 2, self.ndf * 4, 4, 2, 1)),
|
||||
nn.InstanceNorm2d(self.ndf * 4),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
# 16 X 16
|
||||
SpectralNorm(nn.Conv2d(self.ndf * 4, self.ndf * 8, 4, 2, 1)),
|
||||
nn.InstanceNorm2d(self.ndf * 8),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
# 8 X 8
|
||||
SpectralNorm(nn.Conv2d(self.ndf * 8, self.ndf * 16, 4, 2, 1)),
|
||||
nn.InstanceNorm2d(self.ndf * 16),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
)
|
||||
# 4 X 4
|
||||
self.last = SpectralNorm(nn.Conv2d(self.ndf * 16, 1, 4, 1, 0))
|
||||
# self.main = nn.Sequential(
|
||||
# # input size is in_size x 64 x 64
|
||||
# SpectralNorm(nn.Conv2d(self.in_size, self.ndf, 4, 2, 1)),
|
||||
# nn.LeakyReLU(0.2, inplace=True),
|
||||
# # state size: ndf x 32 x 32
|
||||
# Self_Attention(self.ndf),
|
||||
# SpectralNorm(nn.Conv2d(self.ndf, self.ndf * 2, 4, 2, 1)),
|
||||
# nn.InstanceNorm2d(self.ndf * 2),
|
||||
# nn.LeakyReLU(0.2, inplace=True),
|
||||
# # state size: (ndf * 2) x 16 x 16
|
||||
# SpectralNorm(nn.Conv2d(self.ndf * 2, self.ndf * 4, 4, 2, 1)),
|
||||
# nn.InstanceNorm2d(self.ndf * 4),
|
||||
# nn.LeakyReLU(0.2, inplace=True),
|
||||
# # state size: (ndf * 4) x 8 x 8
|
||||
# SpectralNorm(nn.Conv2d(self.ndf * 4, self.ndf * 8, 4, 2, 1)),
|
||||
# nn.InstanceNorm2d(self.ndf * 8),
|
||||
# nn.LeakyReLU(0.2, inplace=True),
|
||||
# # sate size: (ndf * 8) x 4 x 4
|
||||
# )
|
||||
|
||||
def forward(self, input):
|
||||
# self.last = SpectralNorm(nn.Conv2d(self.ndf * 8, 1, 4, 1, 0))
|
||||
|
||||
feature = self.main(input)
|
||||
output = self.last(feature)
|
||||
output = F.avg_pool2d(output, output.size()[2:]).view(output.size()[0], -1)
|
||||
# def forward(self, input):
|
||||
|
||||
return output, feature
|
||||
# feature = self.main(input)
|
||||
# output = self.last(feature)
|
||||
# output = F.avg_pool2d(output, output.size()[2:]).view(output.size()[0], -1)
|
||||
|
||||
# return output, feature
|
||||
|
||||
|
||||
class Discriminator_WITHOUT_FC_x128(nn.Module):
|
||||
"""
|
||||
Discriminative Network
|
||||
"""
|
||||
# class Discriminator_WITHOUT_FC_x64_BIG(nn.Module):
|
||||
# """
|
||||
# Discriminative Network
|
||||
# """
|
||||
|
||||
def __init__(self, in_size=3, ndf=64):
|
||||
super(Discriminator_WITHOUT_FC_x128, self).__init__()
|
||||
self.in_size = in_size
|
||||
self.ndf = ndf
|
||||
# def __init__(self, in_size=3, ndf=64):
|
||||
# super(Discriminator_WITHOUT_FC_x64_BIG, self).__init__()
|
||||
# self.in_size = in_size
|
||||
# self.ndf = ndf
|
||||
|
||||
self.main = nn.Sequential(
|
||||
# input size is in_size x 128 x 128
|
||||
SpectralNorm(nn.Conv2d(self.in_size, self.ndf, 4, 2, 1)),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
# state size: ndf x 64 x 64
|
||||
SpectralNorm(nn.Conv2d(self.ndf, self.ndf, 4, 2, 1)),
|
||||
nn.InstanceNorm2d(self.ndf),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
# state size: ndf x 32 x 32
|
||||
SpectralNorm(nn.Conv2d(self.ndf, self.ndf * 2, 4, 2, 1)),
|
||||
nn.InstanceNorm2d(self.ndf * 2),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
# state size: (ndf * 24) x 16 x 16
|
||||
SpectralNorm(nn.Conv2d(self.ndf * 2, self.ndf * 4, 4, 2, 1)),
|
||||
nn.InstanceNorm2d(self.ndf * 4),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
# sate size: (ndf * 4) x 8 x 8
|
||||
SpectralNorm(nn.Conv2d(self.ndf * 4, self.ndf * 8, 4, 2, 1)),
|
||||
nn.InstanceNorm2d(self.ndf * 8),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
# sate size: (ndf * 8) x 4 x 4
|
||||
)
|
||||
# self.main = nn.Sequential(
|
||||
# # 256 X 256
|
||||
# SpectralNorm(nn.Conv2d(self.in_size, self.ndf, 4, 2, 1)),
|
||||
# nn.LeakyReLU(0.2, inplace=True),
|
||||
# # 128 X 128
|
||||
# SpectralNorm(nn.Conv2d(self.ndf, self.ndf, 4, 2, 1)),
|
||||
# nn.InstanceNorm2d(self.ndf),
|
||||
# nn.LeakyReLU(0.2, inplace=True),
|
||||
# # 64 X 64
|
||||
# Self_Attention(self.ndf),
|
||||
# SpectralNorm(nn.Conv2d(self.ndf, self.ndf * 2, 4, 2, 1)),
|
||||
# nn.InstanceNorm2d(self.ndf * 2),
|
||||
# nn.LeakyReLU(0.2, inplace=True),
|
||||
# # 32 X 32
|
||||
# SpectralNorm(nn.Conv2d(self.ndf * 2, self.ndf * 4, 4, 2, 1)),
|
||||
# nn.InstanceNorm2d(self.ndf * 4),
|
||||
# nn.LeakyReLU(0.2, inplace=True),
|
||||
# # 16 X 16
|
||||
# SpectralNorm(nn.Conv2d(self.ndf * 4, self.ndf * 8, 4, 2, 1)),
|
||||
# nn.InstanceNorm2d(self.ndf * 8),
|
||||
# nn.LeakyReLU(0.2, inplace=True),
|
||||
# # 8 X 8
|
||||
# SpectralNorm(nn.Conv2d(self.ndf * 8, self.ndf * 16, 4, 2, 1)),
|
||||
# nn.InstanceNorm2d(self.ndf * 16),
|
||||
# nn.LeakyReLU(0.2, inplace=True),
|
||||
# )
|
||||
# # 4 X 4
|
||||
# self.last = SpectralNorm(nn.Conv2d(self.ndf * 16, 1, 4, 1, 0))
|
||||
|
||||
self.last = SpectralNorm(nn.Conv2d(self.ndf * 8, 1, 4, 1, 0))
|
||||
# def forward(self, input):
|
||||
|
||||
def forward(self, input):
|
||||
# feature = self.main(input)
|
||||
# output = self.last(feature)
|
||||
# output = F.avg_pool2d(output, output.size()[2:]).view(output.size()[0], -1)
|
||||
|
||||
feature = self.main(input)
|
||||
output = self.last(feature)
|
||||
|
||||
return output, feature
|
||||
# return output, feature
|
||||
|
||||
|
||||
class Discriminator_WITHOUT_FC_x256(nn.Module):
|
||||
"""
|
||||
Discriminative Network
|
||||
"""
|
||||
# class Discriminator_WITHOUT_FC_x128(nn.Module):
|
||||
# """
|
||||
# Discriminative Network
|
||||
# """
|
||||
|
||||
def __init__(self, in_size=3, ndf=64):
|
||||
super(Discriminator_WITHOUT_FC_x256, self).__init__()
|
||||
self.in_size = in_size
|
||||
self.ndf = ndf
|
||||
# def __init__(self, in_size=3, ndf=64):
|
||||
# super(Discriminator_WITHOUT_FC_x128, self).__init__()
|
||||
# self.in_size = in_size
|
||||
# self.ndf = ndf
|
||||
|
||||
self.main = nn.Sequential(
|
||||
# input size is in_size x 256 x 256
|
||||
SpectralNorm(nn.Conv2d(self.in_size, self.ndf, 4, 2, 1)),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
# state size: ndf x 128 x 128
|
||||
SpectralNorm(nn.Conv2d(self.ndf, self.ndf, 4, 2, 1)),
|
||||
nn.InstanceNorm2d(self.ndf),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
# state size: ndf x 64 x 64
|
||||
SpectralNorm(nn.Conv2d(self.ndf, self.ndf, 4, 2, 1)),
|
||||
nn.InstanceNorm2d(self.ndf),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
# state size: (ndf * 1) x 32 x 32
|
||||
SpectralNorm(nn.Conv2d(self.ndf, self.ndf * 2, 4, 2, 1)),
|
||||
nn.InstanceNorm2d(self.ndf * 2),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
# state size: (ndf * 2) x 16 x 16
|
||||
SpectralNorm(nn.Conv2d(self.ndf * 2, self.ndf * 4, 4, 2, 1)),
|
||||
nn.InstanceNorm2d(self.ndf * 4),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
# sate size: (ndf * 4) x 8 x 8
|
||||
SpectralNorm(nn.Conv2d(self.ndf * 4, self.ndf * 8, 4, 2, 1)),
|
||||
nn.InstanceNorm2d(self.ndf * 8),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
# sate size: (ndf * 8) x 4 x 4
|
||||
)
|
||||
# self.main = nn.Sequential(
|
||||
# # input size is in_size x 128 x 128
|
||||
# SpectralNorm(nn.Conv2d(self.in_size, self.ndf, 4, 2, 1)),
|
||||
# nn.LeakyReLU(0.2, inplace=True),
|
||||
# # state size: ndf x 64 x 64
|
||||
# SpectralNorm(nn.Conv2d(self.ndf, self.ndf, 4, 2, 1)),
|
||||
# nn.InstanceNorm2d(self.ndf),
|
||||
# nn.LeakyReLU(0.2, inplace=True),
|
||||
# # state size: ndf x 32 x 32
|
||||
# SpectralNorm(nn.Conv2d(self.ndf, self.ndf * 2, 4, 2, 1)),
|
||||
# nn.InstanceNorm2d(self.ndf * 2),
|
||||
# nn.LeakyReLU(0.2, inplace=True),
|
||||
# # state size: (ndf * 24) x 16 x 16
|
||||
# SpectralNorm(nn.Conv2d(self.ndf * 2, self.ndf * 4, 4, 2, 1)),
|
||||
# nn.InstanceNorm2d(self.ndf * 4),
|
||||
# nn.LeakyReLU(0.2, inplace=True),
|
||||
# # sate size: (ndf * 4) x 8 x 8
|
||||
# SpectralNorm(nn.Conv2d(self.ndf * 4, self.ndf * 8, 4, 2, 1)),
|
||||
# nn.InstanceNorm2d(self.ndf * 8),
|
||||
# nn.LeakyReLU(0.2, inplace=True),
|
||||
# # sate size: (ndf * 8) x 4 x 4
|
||||
# )
|
||||
|
||||
self.last = SpectralNorm(nn.Conv2d(self.ndf * 8, 1, 4, 1, 0))
|
||||
# self.last = SpectralNorm(nn.Conv2d(self.ndf * 8, 1, 4, 1, 0))
|
||||
|
||||
def forward(self, input):
|
||||
# def forward(self, input):
|
||||
|
||||
feature = self.main(input)
|
||||
output = self.last(feature)
|
||||
# feature = self.main(input)
|
||||
# output = self.last(feature)
|
||||
|
||||
return output, feature
|
||||
# return output, feature
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
discriminator1 = Discriminator_WITHOUT_FC_x64()
|
||||
discriminator2 = Discriminator_WITHOUT_FC_x128()
|
||||
discriminator3 = Discriminator_WITHOUT_FC_x256()
|
||||
print(discriminator1)
|
||||
print(discriminator2)
|
||||
print(discriminator3)
|
||||
# class Discriminator_WITHOUT_FC_x256(nn.Module):
|
||||
# """
|
||||
# Discriminative Network
|
||||
# """
|
||||
|
||||
# def __init__(self, in_size=3, ndf=64):
|
||||
# super(Discriminator_WITHOUT_FC_x256, self).__init__()
|
||||
# self.in_size = in_size
|
||||
# self.ndf = ndf
|
||||
|
||||
# self.main = nn.Sequential(
|
||||
# # input size is in_size x 256 x 256
|
||||
# SpectralNorm(nn.Conv2d(self.in_size, self.ndf, 4, 2, 1)),
|
||||
# nn.LeakyReLU(0.2, inplace=True),
|
||||
# # state size: ndf x 128 x 128
|
||||
# SpectralNorm(nn.Conv2d(self.ndf, self.ndf, 4, 2, 1)),
|
||||
# nn.InstanceNorm2d(self.ndf),
|
||||
# nn.LeakyReLU(0.2, inplace=True),
|
||||
# # state size: ndf x 64 x 64
|
||||
# SpectralNorm(nn.Conv2d(self.ndf, self.ndf, 4, 2, 1)),
|
||||
# nn.InstanceNorm2d(self.ndf),
|
||||
# nn.LeakyReLU(0.2, inplace=True),
|
||||
# # state size: (ndf * 1) x 32 x 32
|
||||
# SpectralNorm(nn.Conv2d(self.ndf, self.ndf * 2, 4, 2, 1)),
|
||||
# nn.InstanceNorm2d(self.ndf * 2),
|
||||
# nn.LeakyReLU(0.2, inplace=True),
|
||||
# # state size: (ndf * 2) x 16 x 16
|
||||
# SpectralNorm(nn.Conv2d(self.ndf * 2, self.ndf * 4, 4, 2, 1)),
|
||||
# nn.InstanceNorm2d(self.ndf * 4),
|
||||
# nn.LeakyReLU(0.2, inplace=True),
|
||||
# # sate size: (ndf * 4) x 8 x 8
|
||||
# SpectralNorm(nn.Conv2d(self.ndf * 4, self.ndf * 8, 4, 2, 1)),
|
||||
# nn.InstanceNorm2d(self.ndf * 8),
|
||||
# nn.LeakyReLU(0.2, inplace=True),
|
||||
# # sate size: (ndf * 8) x 4 x 4
|
||||
# )
|
||||
|
||||
# self.last = SpectralNorm(nn.Conv2d(self.ndf * 8, 1, 4, 1, 0))
|
||||
|
||||
# def forward(self, input):
|
||||
|
||||
# feature = self.main(input)
|
||||
# output = self.last(feature)
|
||||
|
||||
# return output, feature
|
||||
|
||||
|
||||
# if __name__ == "__main__":
|
||||
# discriminator1 = Discriminator_WITHOUT_FC_x64()
|
||||
# discriminator2 = Discriminator_WITHOUT_FC_x128()
|
||||
# discriminator3 = Discriminator_WITHOUT_FC_x256()
|
||||
# print(discriminator1)
|
||||
# print(discriminator2)
|
||||
# print(discriminator3)
|
||||
|
||||
@@ -189,13 +189,13 @@ class Self_Attn(nn.Module):
|
||||
return out
|
||||
|
||||
|
||||
class VGG19_feature_color_torchversion(nn.Module):
|
||||
class VGG19_pytorch(nn.Module):
|
||||
"""
|
||||
NOTE: no need to pre-process the input; input tensor should range in [0,1]
|
||||
"""
|
||||
|
||||
def __init__(self, pool="max"):
|
||||
super(VGG19_feature_color_torchversion, self).__init__()
|
||||
super(VGG19_pytorch, self).__init__()
|
||||
self.conv1_1 = nn.Conv2d(3, 64, kernel_size=3, padding=1)
|
||||
self.conv1_2 = nn.Conv2d(64, 64, kernel_size=3, padding=1)
|
||||
self.conv2_1 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
|
||||
|
||||
Reference in New Issue
Block a user