197 lines
7.9 KiB
Python
197 lines
7.9 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import functools
|
|
from model.registry import MODEL
|
|
|
|
|
|
def _select_norm_layer(norm_type):
|
|
if norm_type == "BN":
|
|
return functools.partial(nn.BatchNorm2d, affine=True, track_running_stats=True)
|
|
elif norm_type == "IN":
|
|
return functools.partial(nn.InstanceNorm2d, affine=False, track_running_stats=False)
|
|
elif norm_type == "NONE":
|
|
return lambda x: nn.Identity()
|
|
else:
|
|
raise NotImplemented(f'normalization layer {norm_type} is not found')
|
|
|
|
|
|
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
|
|
|
|
|
|
@MODEL.register_module()
|
|
class ResidualBlock(nn.Module):
|
|
def __init__(self, num_channels, padding_mode='reflect', norm_type="IN", use_dropout=False, 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)
|
|
models = [nn.Sequential(
|
|
nn.Conv2d(num_channels, num_channels, kernel_size=3, padding=1, padding_mode=padding_mode, bias=use_bias),
|
|
norm_layer(num_channels),
|
|
nn.ReLU(inplace=True),
|
|
)]
|
|
if use_dropout:
|
|
models.append(nn.Dropout(0.5))
|
|
models.append(nn.Sequential(
|
|
nn.Conv2d(num_channels, num_channels, kernel_size=3, padding=1, padding_mode=padding_mode, bias=use_bias),
|
|
norm_layer(num_channels),
|
|
))
|
|
self.block = nn.Sequential(*models)
|
|
|
|
def forward(self, x):
|
|
return x + self.block(x)
|
|
|
|
|
|
@MODEL.register_module()
|
|
class ResGenerator(nn.Module):
|
|
def __init__(self, in_channels, out_channels, base_channels=64, num_blocks=9, padding_mode='reflect',
|
|
norm_type="IN", use_dropout=False):
|
|
super(ResGenerator, self).__init__()
|
|
assert num_blocks >= 0, f'Number of residual blocks must be non-negative, but got {num_blocks}.'
|
|
norm_layer = _select_norm_layer(norm_type)
|
|
use_bias = norm_type == "IN"
|
|
|
|
self.start_conv = nn.Sequential(
|
|
nn.Conv2d(in_channels, base_channels, kernel_size=7, stride=1, padding_mode=padding_mode, padding=3,
|
|
bias=use_bias),
|
|
norm_layer(num_features=base_channels),
|
|
nn.ReLU(inplace=True)
|
|
)
|
|
|
|
# down sampling
|
|
submodules = []
|
|
num_down_sampling = 2
|
|
for i in range(num_down_sampling):
|
|
multiple = 2 ** i
|
|
submodules += [
|
|
nn.Conv2d(in_channels=base_channels * multiple, out_channels=base_channels * multiple * 2,
|
|
kernel_size=3,
|
|
stride=2, padding=1, bias=use_bias),
|
|
norm_layer(num_features=base_channels * multiple * 2),
|
|
nn.ReLU(inplace=True)
|
|
]
|
|
self.encoder = nn.Sequential(*submodules)
|
|
|
|
res_block_channels = num_down_sampling ** 2 * base_channels
|
|
self.res_blocks = nn.ModuleList(
|
|
[ResidualBlock(res_block_channels, padding_mode, norm_type, use_dropout=use_dropout) for _ in
|
|
range(num_blocks)])
|
|
|
|
# up sampling
|
|
submodules = []
|
|
for i in range(num_down_sampling):
|
|
multiple = 2 ** (num_down_sampling - i)
|
|
submodules += [
|
|
nn.ConvTranspose2d(base_channels * multiple, base_channels * multiple // 2, kernel_size=3, stride=2,
|
|
padding=1, output_padding=1, bias=use_bias),
|
|
norm_layer(num_features=base_channels * multiple // 2),
|
|
nn.ReLU(inplace=True),
|
|
]
|
|
self.decoder = nn.Sequential(*submodules)
|
|
|
|
self.end_conv = nn.Sequential(
|
|
nn.Conv2d(base_channels, out_channels, kernel_size=7, padding=3, padding_mode=padding_mode),
|
|
nn.Tanh()
|
|
)
|
|
|
|
def forward(self, x):
|
|
x = self.encoder(self.start_conv(x))
|
|
for rb in self.res_blocks:
|
|
x = rb(x)
|
|
return self.end_conv(self.decoder(x))
|
|
|
|
|
|
@MODEL.register_module()
|
|
class PatchDiscriminator(nn.Module):
|
|
def __init__(self, in_channels, base_channels=64, num_conv=3, norm_type="BN"):
|
|
super(PatchDiscriminator, self).__init__()
|
|
assert num_conv >= 0, f'Number of conv blocks must be non-negative, but got {num_conv}.'
|
|
norm_layer = _select_norm_layer(norm_type)
|
|
use_bias = norm_type == "IN"
|
|
|
|
kernel_size = 4
|
|
padding = 1
|
|
sequence = [
|
|
nn.Conv2d(in_channels, base_channels, kernel_size=kernel_size, stride=2, padding=padding),
|
|
nn.LeakyReLU(0.2, inplace=True),
|
|
]
|
|
|
|
# stacked intermediate layers,
|
|
# gradually increasing the number of filters
|
|
multiple_now = 1
|
|
for n in range(1, num_conv):
|
|
multiple_prev = multiple_now
|
|
multiple_now = min(2 ** n, 8)
|
|
sequence += [
|
|
nn.Conv2d(base_channels * multiple_prev, base_channels * multiple_now, kernel_size=kernel_size,
|
|
padding=padding, stride=2, bias=use_bias),
|
|
norm_layer(base_channels * multiple_now),
|
|
nn.LeakyReLU(0.2, inplace=True)
|
|
]
|
|
multiple_prev = multiple_now
|
|
multiple_now = min(2 ** num_conv, 8)
|
|
sequence += [
|
|
nn.Conv2d(base_channels * multiple_prev, base_channels * multiple_now, kernel_size, stride=1,
|
|
padding=padding, bias=use_bias),
|
|
norm_layer(base_channels * multiple_now),
|
|
nn.LeakyReLU(0.2, inplace=True),
|
|
nn.Conv2d(base_channels * multiple_now, 1, kernel_size, stride=1, padding=padding)
|
|
]
|
|
self.model = nn.Sequential(*sequence)
|
|
|
|
def forward(self, x):
|
|
return self.model(x)
|