204 lines
8.1 KiB
Python
204 lines
8.1 KiB
Python
from functools import partial
|
|
|
|
import math
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
from model import MODEL
|
|
from model.normalization import select_norm_layer
|
|
|
|
|
|
class GANImageBuffer(object):
|
|
"""This class implements an image buffer that stores previously
|
|
generated images.
|
|
This buffer allows us to update the discriminator using a history of
|
|
generated images rather than the ones produced by the latest generator
|
|
to reduce model oscillation.
|
|
Args:
|
|
buffer_size (int): The size of image buffer. If buffer_size = 0,
|
|
no buffer will be created.
|
|
buffer_ratio (float): The chance / possibility to use the images
|
|
previously stored in the buffer.
|
|
"""
|
|
|
|
def __init__(self, buffer_size, buffer_ratio=0.5):
|
|
self.buffer_size = buffer_size
|
|
# create an empty buffer
|
|
if self.buffer_size > 0:
|
|
self.img_num = 0
|
|
self.image_buffer = []
|
|
self.buffer_ratio = buffer_ratio
|
|
|
|
def query(self, images):
|
|
"""Query current image batch using a history of generated images.
|
|
Args:
|
|
images (Tensor): Current image batch without history information.
|
|
"""
|
|
if self.buffer_size == 0: # if the buffer size is 0, do nothing
|
|
return images
|
|
return_images = []
|
|
for image in images:
|
|
image = torch.unsqueeze(image.data, 0)
|
|
# if the buffer is not full, keep inserting current images
|
|
if self.img_num < self.buffer_size:
|
|
self.img_num = self.img_num + 1
|
|
self.image_buffer.append(image)
|
|
return_images.append(image)
|
|
else:
|
|
use_buffer = torch.rand(1) < self.buffer_ratio
|
|
# by self.buffer_ratio, the buffer will return a previously
|
|
# stored image, and insert the current image into the buffer
|
|
if use_buffer:
|
|
random_id = torch.randint(0, self.buffer_size, (1,)).item()
|
|
image_tmp = self.image_buffer[random_id].clone()
|
|
self.image_buffer[random_id] = image
|
|
return_images.append(image_tmp)
|
|
# by (1 - self.buffer_ratio), the buffer will return the
|
|
# current image
|
|
else:
|
|
return_images.append(image)
|
|
# collect all the images and return
|
|
return_images = torch.cat(return_images, 0)
|
|
return return_images
|
|
|
|
|
|
# based SPADE or pix2pixHD Discriminator
|
|
@MODEL.register_module("PatchDiscriminator")
|
|
class PatchDiscriminator(nn.Module):
|
|
def __init__(self, in_channels, base_channels, num_conv=4, use_spectral=False, norm_type="IN",
|
|
need_intermediate_feature=False):
|
|
super().__init__()
|
|
self.need_intermediate_feature = need_intermediate_feature
|
|
|
|
kernel_size = 4
|
|
padding = math.ceil((kernel_size - 1.0) / 2)
|
|
norm_layer = select_norm_layer(norm_type)
|
|
use_bias = norm_type == "IN"
|
|
padding_mode = "zeros"
|
|
|
|
sequence = [nn.Sequential(
|
|
nn.Conv2d(in_channels, base_channels, kernel_size, stride=2, padding=padding),
|
|
nn.LeakyReLU(0.2, False)
|
|
)]
|
|
multiple_now = 1
|
|
for i in range(1, num_conv):
|
|
multiple_prev = multiple_now
|
|
multiple_now = min(2 ** i, 2 ** 3)
|
|
stride = 1 if i == num_conv - 1 else 2
|
|
sequence.append(nn.Sequential(
|
|
self.build_conv2d(use_spectral, base_channels * multiple_prev, base_channels * multiple_now,
|
|
kernel_size, stride, padding, bias=use_bias, padding_mode=padding_mode),
|
|
norm_layer(base_channels * multiple_now),
|
|
nn.LeakyReLU(0.2, inplace=False),
|
|
))
|
|
multiple_now = min(2 ** num_conv, 8)
|
|
sequence.append(nn.Conv2d(base_channels * multiple_now, 1, kernel_size, stride=1, padding=padding,
|
|
padding_mode=padding_mode))
|
|
self.conv_blocks = nn.ModuleList(sequence)
|
|
|
|
@staticmethod
|
|
def build_conv2d(use_spectral, in_channels: int, out_channels: int, kernel_size, stride, padding,
|
|
bias=True, padding_mode: str = 'zeros'):
|
|
conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, bias, padding_mode=padding_mode)
|
|
if not use_spectral:
|
|
return conv
|
|
return nn.utils.spectral_norm(conv)
|
|
|
|
def forward(self, x):
|
|
if self.need_intermediate_feature:
|
|
intermediate_feature = []
|
|
for layer in self.conv_blocks:
|
|
x = layer(x)
|
|
intermediate_feature.append(x)
|
|
return tuple(intermediate_feature)
|
|
else:
|
|
for layer in self.conv_blocks:
|
|
x = layer(x)
|
|
return x
|
|
|
|
|
|
@MODEL.register_module()
|
|
class ResidualBlock(nn.Module):
|
|
def __init__(self, num_channels, padding_mode='reflect', norm_type="IN", use_bias=None):
|
|
super(ResidualBlock, self).__init__()
|
|
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)
|
|
self.conv1 = nn.Conv2d(num_channels, num_channels, kernel_size=3, padding=1, padding_mode=padding_mode,
|
|
bias=use_bias)
|
|
self.norm1 = norm_layer(num_channels)
|
|
self.relu1 = nn.ReLU(inplace=True)
|
|
self.conv2 = nn.Conv2d(num_channels, num_channels, kernel_size=3, padding=1, padding_mode=padding_mode,
|
|
bias=use_bias)
|
|
self.norm2 = norm_layer(num_channels)
|
|
|
|
def forward(self, x):
|
|
res = x
|
|
x = self.relu1(self.norm1(self.conv1(x)))
|
|
x = self.norm2(self.conv2(x))
|
|
return x + res
|
|
|
|
|
|
_DO_NO_THING_FUNC = lambda x: x
|
|
|
|
|
|
def select_activation(t):
|
|
if t == "ReLU":
|
|
return partial(nn.ReLU, inplace=True)
|
|
elif t == "LeakyReLU":
|
|
return partial(nn.LeakyReLU, negative_slope=0.2, inplace=True)
|
|
elif t == "Tanh":
|
|
return partial(nn.Tanh)
|
|
elif t == "NONE":
|
|
return _DO_NO_THING_FUNC
|
|
else:
|
|
raise NotImplemented
|
|
|
|
|
|
def _use_bias_checker(norm_type):
|
|
return norm_type not in ["IN", "BN", "AdaIN"]
|
|
|
|
|
|
class Conv2dBlock(nn.Module):
|
|
def __init__(self, in_channels: int, out_channels: int, use_spectral_norm=False, activation_type="ReLU",
|
|
bias=None, norm_type="NONE", **conv_kwargs):
|
|
super().__init__()
|
|
self.norm_type = norm_type
|
|
self.activation_type = activation_type
|
|
conv_kwargs["bias"] = _use_bias_checker(norm_type) if bias is None else bias
|
|
conv = nn.Conv2d(in_channels, out_channels, **conv_kwargs)
|
|
self.convolution = nn.utils.spectral_norm(conv) if use_spectral_norm else conv
|
|
if norm_type != "NONE":
|
|
self.normalization = select_norm_layer(norm_type)(out_channels)
|
|
if activation_type != "NONE":
|
|
self.activation = select_activation(activation_type)()
|
|
|
|
def forward(self, x):
|
|
x = self.convolution(x)
|
|
if self.norm_type != "NONE":
|
|
x = self.normalization(x)
|
|
if self.activation_type != "NONE":
|
|
x = self.activation(x)
|
|
return x
|
|
|
|
|
|
class ResBlock(nn.Module):
|
|
def __init__(self, num_channels, use_spectral_norm=False, padding_mode='reflect',
|
|
norm_type="IN", activation_type="ReLU", use_bias=None):
|
|
super().__init__()
|
|
self.norm_type = norm_type
|
|
if use_bias is None:
|
|
# bias will be canceled after channel wise normalization
|
|
use_bias = _use_bias_checker(norm_type)
|
|
|
|
self.conv1 = Conv2dBlock(num_channels, num_channels, use_spectral_norm,
|
|
kernel_size=3, padding=1, padding_mode=padding_mode, bias=use_bias,
|
|
norm_type=norm_type, activation_type=activation_type)
|
|
self.conv2 = Conv2dBlock(num_channels, num_channels, use_spectral_norm,
|
|
kernel_size=3, padding=1, padding_mode=padding_mode, bias=use_bias,
|
|
norm_type=norm_type, activation_type="NONE")
|
|
|
|
def forward(self, x):
|
|
return self.conv2(self.conv1(x)) + x
|