add U-GAT-IT
This commit is contained in:
253
model/GAN/UGATIT.py
Normal file
253
model/GAN/UGATIT.py
Normal file
@@ -0,0 +1,253 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from .residual_generator import ResidualBlock
|
||||
from model.registry import MODEL
|
||||
|
||||
|
||||
class RhoClipper(object):
|
||||
def __init__(self, clip_min, clip_max):
|
||||
self.clip_min = clip_min
|
||||
self.clip_max = clip_max
|
||||
assert clip_min < clip_max
|
||||
|
||||
def __call__(self, module):
|
||||
if hasattr(module, 'rho'):
|
||||
w = module.rho.data
|
||||
w = w.clamp(self.clip_min, self.clip_max)
|
||||
module.rho.data = w
|
||||
|
||||
|
||||
@MODEL.register_module("UGATIT-Generator")
|
||||
class Generator(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, base_channels=64, num_blocks=6, img_size=256, light=False):
|
||||
assert (num_blocks >= 0)
|
||||
super(Generator, self).__init__()
|
||||
self.input_channels = in_channels
|
||||
self.output_channels = out_channels
|
||||
self.base_channels = base_channels
|
||||
self.num_blocks = num_blocks
|
||||
self.img_size = img_size
|
||||
self.light = light
|
||||
|
||||
down_encoder = [nn.Conv2d(in_channels, base_channels, kernel_size=7, stride=1, padding=3,
|
||||
padding_mode="reflect", bias=False),
|
||||
nn.InstanceNorm2d(base_channels),
|
||||
nn.ReLU(True)]
|
||||
|
||||
n_down_sampling = 2
|
||||
for i in range(n_down_sampling):
|
||||
mult = 2 ** i
|
||||
down_encoder += [nn.Conv2d(base_channels * mult, base_channels * mult * 2, kernel_size=3, stride=2,
|
||||
padding=1, bias=False, padding_mode="reflect"),
|
||||
nn.InstanceNorm2d(base_channels * mult * 2),
|
||||
nn.ReLU(True)]
|
||||
|
||||
# Down-Sampling Bottleneck
|
||||
mult = 2 ** n_down_sampling
|
||||
for i in range(num_blocks):
|
||||
# TODO: change ResnetBlock to ResidualBlock, check use_bias param
|
||||
down_encoder += [ResidualBlock(base_channels * mult, use_bias=False)]
|
||||
self.down_encoder = nn.Sequential(*down_encoder)
|
||||
|
||||
# Class Activation Map
|
||||
self.gap_fc = nn.Linear(base_channels * mult, 1, bias=False)
|
||||
self.gmp_fc = nn.Linear(base_channels * mult, 1, bias=False)
|
||||
self.conv1x1 = nn.Conv2d(base_channels * mult * 2, base_channels * mult, kernel_size=1, stride=1, bias=True)
|
||||
self.relu = nn.ReLU(True)
|
||||
|
||||
# Gamma, Beta block
|
||||
if self.light:
|
||||
fc = [nn.Linear(base_channels * mult, base_channels * mult, bias=False),
|
||||
nn.ReLU(True),
|
||||
nn.Linear(base_channels * mult, base_channels * mult, bias=False),
|
||||
nn.ReLU(True)]
|
||||
else:
|
||||
fc = [
|
||||
nn.Linear(img_size // mult * img_size // mult * base_channels * mult, base_channels * mult, bias=False),
|
||||
nn.ReLU(True),
|
||||
nn.Linear(base_channels * mult, base_channels * mult, bias=False),
|
||||
nn.ReLU(True)]
|
||||
self.fc = nn.Sequential(*fc)
|
||||
|
||||
self.gamma = nn.Linear(base_channels * mult, base_channels * mult, bias=False)
|
||||
self.beta = nn.Linear(base_channels * mult, base_channels * mult, bias=False)
|
||||
|
||||
# Up-Sampling Bottleneck
|
||||
self.up_bottleneck = nn.ModuleList(
|
||||
[ResnetAdaILNBlock(base_channels * mult, use_bias=False) for _ in range(num_blocks)])
|
||||
|
||||
# Up-Sampling
|
||||
up_decoder = []
|
||||
for i in range(n_down_sampling):
|
||||
mult = 2 ** (n_down_sampling - i)
|
||||
up_decoder += [nn.Upsample(scale_factor=2, mode='nearest'),
|
||||
nn.Conv2d(base_channels * mult, base_channels * mult // 2, kernel_size=3, stride=1,
|
||||
padding=1, padding_mode="reflect", bias=False),
|
||||
ILN(base_channels * mult // 2),
|
||||
nn.ReLU(True)]
|
||||
|
||||
up_decoder += [nn.Conv2d(base_channels, out_channels, kernel_size=7, stride=1, padding=3,
|
||||
padding_mode="reflect", bias=False),
|
||||
nn.Tanh()]
|
||||
self.up_decoder = nn.Sequential(*up_decoder)
|
||||
# self.up_decoder = nn.ModuleDict({
|
||||
# "up_1": nn.Upsample(scale_factor=2, mode='nearest'),
|
||||
# "up_conv_1": nn.Sequential(
|
||||
# nn.Conv2d(base_channels * 4, base_channels * 4 // 2, kernel_size=3, stride=1,
|
||||
# padding=1, padding_mode="reflect", bias=False),
|
||||
# ILN(base_channels * 4 // 2),
|
||||
# nn.ReLU(True)),
|
||||
# "up_2": nn.Upsample(scale_factor=2, mode='nearest'),
|
||||
# "up_conv_2": nn.Sequential(
|
||||
# nn.Conv2d(base_channels * 2, base_channels * 2 // 2, kernel_size=3, stride=1,
|
||||
# padding=1, padding_mode="reflect", bias=False),
|
||||
# ILN(base_channels * 2 // 2),
|
||||
# nn.ReLU(True)),
|
||||
# "up_end": nn.Sequential(nn.Conv2d(base_channels, out_channels, kernel_size=7, stride=1, padding=3,
|
||||
# padding_mode="reflect", bias=False), nn.Tanh())
|
||||
# })
|
||||
|
||||
def forward(self, x):
|
||||
x = self.down_encoder(x)
|
||||
|
||||
gap = torch.nn.functional.adaptive_avg_pool2d(x, 1)
|
||||
gap_logit = self.gap_fc(gap.view(x.shape[0], -1))
|
||||
gap = x * self.gap_fc.weight.unsqueeze(2).unsqueeze(3)
|
||||
|
||||
gmp = torch.nn.functional.adaptive_max_pool2d(x, 1)
|
||||
gmp_logit = self.gmp_fc(gmp.view(x.shape[0], -1))
|
||||
gmp = x * self.gmp_fc.weight.unsqueeze(2).unsqueeze(3)
|
||||
|
||||
cam_logit = torch.cat([gap_logit, gmp_logit], 1)
|
||||
|
||||
x = torch.cat([gap, gmp], 1)
|
||||
x = self.relu(self.conv1x1(x))
|
||||
|
||||
heatmap = torch.sum(x, dim=1, keepdim=True)
|
||||
|
||||
if self.light:
|
||||
x_ = torch.nn.functional.adaptive_avg_pool2d(x, 1)
|
||||
x_ = self.fc(x_.view(x_.shape[0], -1))
|
||||
else:
|
||||
x_ = self.fc(x.view(x.shape[0], -1))
|
||||
gamma, beta = self.gamma(x_), self.beta(x_)
|
||||
|
||||
for ub in self.up_bottleneck:
|
||||
x = ub(x, gamma, beta)
|
||||
|
||||
x = self.up_decoder(x)
|
||||
return x, cam_logit, heatmap
|
||||
|
||||
|
||||
class ResnetAdaILNBlock(nn.Module):
|
||||
def __init__(self, dim, use_bias):
|
||||
super(ResnetAdaILNBlock, self).__init__()
|
||||
self.conv1 = nn.Conv2d(dim, dim, kernel_size=3, stride=1, padding=1, bias=use_bias, padding_mode="reflect")
|
||||
self.norm1 = AdaILN(dim)
|
||||
self.relu1 = nn.ReLU(True)
|
||||
|
||||
self.conv2 = nn.Conv2d(dim, dim, kernel_size=3, stride=1, padding=1, bias=use_bias, padding_mode="reflect")
|
||||
self.norm2 = AdaILN(dim)
|
||||
|
||||
def forward(self, x, gamma, beta):
|
||||
out = self.conv1(x)
|
||||
out = self.norm1(out, gamma, beta)
|
||||
out = self.relu1(out)
|
||||
out = self.conv2(out)
|
||||
out = self.norm2(out, gamma, beta)
|
||||
|
||||
return out + x
|
||||
|
||||
|
||||
def instance_layer_normalization(x, gamma, beta, rho, eps=1e-5):
|
||||
in_mean, in_var = torch.mean(x, dim=[2, 3], keepdim=True), torch.var(x, dim=[2, 3], keepdim=True)
|
||||
out_in = (x - in_mean) / torch.sqrt(in_var + eps)
|
||||
ln_mean, ln_var = torch.mean(x, dim=[1, 2, 3], keepdim=True), torch.var(x, dim=[1, 2, 3], keepdim=True)
|
||||
out_ln = (x - ln_mean) / torch.sqrt(ln_var + eps)
|
||||
out = rho.expand(x.shape[0], -1, -1, -1) * out_in + (1 - rho.expand(x.shape[0], -1, -1, -1)) * out_ln
|
||||
out = out * gamma.unsqueeze(2).unsqueeze(3) + beta.unsqueeze(2).unsqueeze(3)
|
||||
return out
|
||||
|
||||
|
||||
class AdaILN(nn.Module):
|
||||
def __init__(self, num_features, eps=1e-5, default_rho=0.9):
|
||||
super(AdaILN, self).__init__()
|
||||
self.eps = eps
|
||||
self.rho = nn.Parameter(torch.Tensor(1, num_features, 1, 1))
|
||||
self.rho.data.fill_(default_rho)
|
||||
|
||||
def forward(self, x, gamma, beta):
|
||||
return instance_layer_normalization(x, gamma, beta, self.rho, self.eps)
|
||||
|
||||
|
||||
class ILN(nn.Module):
|
||||
def __init__(self, num_features, eps=1e-5):
|
||||
super(ILN, self).__init__()
|
||||
self.eps = eps
|
||||
self.rho = nn.Parameter(torch.Tensor(1, num_features, 1, 1))
|
||||
self.gamma = nn.Parameter(torch.Tensor(1, num_features))
|
||||
self.beta = nn.Parameter(torch.Tensor(1, num_features))
|
||||
self.rho.data.fill_(0.0)
|
||||
self.gamma.data.fill_(1.0)
|
||||
self.beta.data.fill_(0.0)
|
||||
|
||||
def forward(self, x):
|
||||
return instance_layer_normalization(
|
||||
x, self.gamma.expand(x.shape[0], -1), self.beta.expand(x.shape[0], -1), self.rho, self.eps)
|
||||
|
||||
|
||||
@MODEL.register_module("UGATIT-Discriminator")
|
||||
class Discriminator(nn.Module):
|
||||
def __init__(self, in_channels, base_channels=64, num_blocks=5):
|
||||
super(Discriminator, self).__init__()
|
||||
encoder = [self.build_conv_block(in_channels, base_channels)]
|
||||
|
||||
for i in range(1, num_blocks - 2):
|
||||
mult = 2 ** (i - 1)
|
||||
encoder.append(self.build_conv_block(base_channels * mult, base_channels * mult * 2))
|
||||
|
||||
mult = 2 ** (num_blocks - 2 - 1)
|
||||
encoder.append(self.build_conv_block(base_channels * mult, base_channels * mult * 2, stride=1))
|
||||
|
||||
self.encoder = nn.Sequential(*encoder)
|
||||
|
||||
# Class Activation Map
|
||||
mult = 2 ** (num_blocks - 2)
|
||||
self.gap_fc = nn.utils.spectral_norm(nn.Linear(base_channels * mult, 1, bias=False))
|
||||
self.gmp_fc = nn.utils.spectral_norm(nn.Linear(base_channels * mult, 1, bias=False))
|
||||
self.conv1x1 = nn.Conv2d(base_channels * mult * 2, base_channels * mult, kernel_size=1, stride=1, bias=True)
|
||||
self.leaky_relu = nn.LeakyReLU(0.2, True)
|
||||
|
||||
self.conv = nn.utils.spectral_norm(
|
||||
nn.Conv2d(base_channels * mult, 1, kernel_size=4, stride=1, padding=1, bias=False, padding_mode="reflect"))
|
||||
|
||||
@staticmethod
|
||||
def build_conv_block(in_channels, out_channels, kernel_size=4, stride=2, padding=1, padding_mode="reflect"):
|
||||
return nn.Sequential(*[
|
||||
nn.utils.spectral_norm(nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, stride=stride,
|
||||
bias=True, padding=padding, padding_mode=padding_mode)),
|
||||
nn.LeakyReLU(0.2, True),
|
||||
])
|
||||
|
||||
def forward(self, x, return_heatmap=False):
|
||||
x = self.encoder(x)
|
||||
batch_size = x.size(0)
|
||||
|
||||
gap = torch.nn.functional.adaptive_avg_pool2d(x, 1) # B x C x 1 x 1, avg of per channel
|
||||
gap_logit = self.gap_fc(gap.view(batch_size, -1))
|
||||
gap = x * self.gap_fc.weight.unsqueeze(2).unsqueeze(3)
|
||||
|
||||
gmp = torch.nn.functional.adaptive_max_pool2d(x, 1)
|
||||
gmp_logit = self.gmp_fc(gmp.view(batch_size, -1))
|
||||
gmp = x * self.gmp_fc.weight.unsqueeze(2).unsqueeze(3)
|
||||
|
||||
cam_logit = torch.cat([gap_logit, gmp_logit], 1)
|
||||
|
||||
x = torch.cat([gap, gmp], 1)
|
||||
x = self.leaky_relu(self.conv1x1(x))
|
||||
|
||||
if return_heatmap:
|
||||
heatmap = torch.sum(x, dim=1, keepdim=True)
|
||||
return self.conv(x), cam_logit, heatmap
|
||||
else:
|
||||
return self.conv(x), cam_logit
|
||||
0
model/GAN/__init__.py
Normal file
0
model/GAN/__init__.py
Normal file
@@ -1,7 +1,7 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import functools
|
||||
from .registry import MODEL
|
||||
from model.registry import MODEL
|
||||
|
||||
|
||||
def _select_norm_layer(norm_type):
|
||||
@@ -71,11 +71,12 @@ class GANImageBuffer(object):
|
||||
|
||||
@MODEL.register_module()
|
||||
class ResidualBlock(nn.Module):
|
||||
def __init__(self, num_channels, padding_mode='reflect', norm_type="IN", use_dropout=False):
|
||||
def __init__(self, num_channels, padding_mode='reflect', norm_type="IN", use_dropout=False, use_bias=None):
|
||||
super(ResidualBlock, self).__init__()
|
||||
|
||||
# Only for IN, use bias since it does not have affine parameters.
|
||||
use_bias = norm_type == "IN"
|
||||
if use_bias is None:
|
||||
# Only for IN, use bias since it does not have affine parameters.
|
||||
use_bias = norm_type == "IN"
|
||||
norm_layer = _select_norm_layer(norm_type)
|
||||
models = [nn.Sequential(
|
||||
nn.Conv2d(num_channels, num_channels, kernel_size=3, padding=1, padding_mode=padding_mode, bias=use_bias),
|
||||
@@ -1,3 +1,3 @@
|
||||
from model.registry import MODEL
|
||||
import model.residual_generator
|
||||
import model.GAN.residual_generator
|
||||
import model.fewshot
|
||||
|
||||
Reference in New Issue
Block a user