|
|
import math
|
|
|
|
|
|
import torch
|
|
|
import torch.nn as nn
|
|
|
|
|
|
def default_conv(in_channels, out_channels, kernel_size, bias=True, groups=1):
|
|
|
return nn.Conv2d(
|
|
|
in_channels, out_channels, kernel_size,
|
|
|
padding=(kernel_size // 2), bias=bias, groups=groups)
|
|
|
|
|
|
def default_norm(n_feats):
|
|
|
return nn.BatchNorm2d(n_feats)
|
|
|
|
|
|
def default_act():
|
|
|
return nn.ReLU(True)
|
|
|
|
|
|
def empty_h(x, n_feats):
|
|
|
'''
|
|
|
create an empty hidden state
|
|
|
|
|
|
input
|
|
|
x: B x T x 3 x H x W
|
|
|
|
|
|
output
|
|
|
h: B x C x H/4 x W/4
|
|
|
'''
|
|
|
b = x.size(0)
|
|
|
h, w = x.size()[-2:]
|
|
|
return x.new_zeros((b, n_feats, h//4, w//4))
|
|
|
|
|
|
class Normalization(nn.Conv2d):
|
|
|
"""Normalize input tensor value with convolutional layer"""
|
|
|
def __init__(self, mean=(0, 0, 0), std=(1, 1, 1)):
|
|
|
super(Normalization, self).__init__(3, 3, kernel_size=1)
|
|
|
tensor_mean = torch.Tensor(mean)
|
|
|
tensor_inv_std = torch.Tensor(std).reciprocal()
|
|
|
|
|
|
self.weight.data = torch.eye(3).mul(tensor_inv_std).view(3, 3, 1, 1)
|
|
|
self.bias.data = torch.Tensor(-tensor_mean.mul(tensor_inv_std))
|
|
|
|
|
|
for params in self.parameters():
|
|
|
params.requires_grad = False
|
|
|
|
|
|
class BasicBlock(nn.Sequential):
|
|
|
"""Convolution layer + Activation layer"""
|
|
|
def __init__(
|
|
|
self, in_channels, out_channels, kernel_size, bias=True,
|
|
|
conv=default_conv, norm=False, act=default_act):
|
|
|
|
|
|
modules = []
|
|
|
modules.append(
|
|
|
conv(in_channels, out_channels, kernel_size, bias=bias))
|
|
|
if norm: modules.append(norm(out_channels))
|
|
|
if act: modules.append(act())
|
|
|
|
|
|
super(BasicBlock, self).__init__(*modules)
|
|
|
|
|
|
class ResBlock(nn.Module):
|
|
|
def __init__(
|
|
|
self, n_feats, kernel_size, bias=True,
|
|
|
conv=default_conv, norm=False, act=default_act):
|
|
|
|
|
|
super(ResBlock, self).__init__()
|
|
|
|
|
|
modules = []
|
|
|
for i in range(2):
|
|
|
modules.append(conv(n_feats, n_feats, kernel_size, bias=bias))
|
|
|
if norm: modules.append(norm(n_feats))
|
|
|
if act and i == 0: modules.append(act())
|
|
|
|
|
|
self.body = nn.Sequential(*modules)
|
|
|
|
|
|
def forward(self, x):
|
|
|
res = self.body(x)
|
|
|
res += x
|
|
|
|
|
|
return res
|
|
|
|
|
|
class ResBlock_mobile(nn.Module):
|
|
|
def __init__(
|
|
|
self, n_feats, kernel_size, bias=True,
|
|
|
conv=default_conv, norm=False, act=default_act, dropout=False):
|
|
|
|
|
|
super(ResBlock_mobile, self).__init__()
|
|
|
|
|
|
modules = []
|
|
|
for i in range(2):
|
|
|
modules.append(conv(n_feats, n_feats, kernel_size, bias=False, groups=n_feats))
|
|
|
modules.append(conv(n_feats, n_feats, 1, bias=False))
|
|
|
if dropout and i == 0: modules.append(nn.Dropout2d(dropout))
|
|
|
if norm: modules.append(norm(n_feats))
|
|
|
if act and i == 0: modules.append(act())
|
|
|
|
|
|
self.body = nn.Sequential(*modules)
|
|
|
|
|
|
def forward(self, x):
|
|
|
res = self.body(x)
|
|
|
res += x
|
|
|
|
|
|
return res
|
|
|
|
|
|
class Upsampler(nn.Sequential):
|
|
|
def __init__(
|
|
|
self, scale, n_feats, bias=True,
|
|
|
conv=default_conv, norm=False, act=False):
|
|
|
|
|
|
modules = []
|
|
|
if (scale & (scale - 1)) == 0:
|
|
|
for _ in range(int(math.log(scale, 2))):
|
|
|
modules.append(conv(n_feats, 4 * n_feats, 3, bias))
|
|
|
modules.append(nn.PixelShuffle(2))
|
|
|
if norm: modules.append(norm(n_feats))
|
|
|
if act: modules.append(act())
|
|
|
elif scale == 3:
|
|
|
modules.append(conv(n_feats, 9 * n_feats, 3, bias))
|
|
|
modules.append(nn.PixelShuffle(3))
|
|
|
if norm: modules.append(norm(n_feats))
|
|
|
if act: modules.append(act())
|
|
|
else:
|
|
|
raise NotImplementedError
|
|
|
|
|
|
super(Upsampler, self).__init__(*modules)
|
|
|
|
|
|
|
|
|
class PixelSort(nn.Module):
|
|
|
"""The inverse operation of PixelShuffle
|
|
|
Reduces the spatial resolution, increasing the number of channels.
|
|
|
Currently, scale 0.5 is supported only.
|
|
|
Later, torch.nn.functional.pixel_sort may be implemented.
|
|
|
Reference:
|
|
|
http://pytorch.org/docs/0.3.0/_modules/torch/nn/modules/pixelshuffle.html#PixelShuffle
|
|
|
http://pytorch.org/docs/0.3.0/_modules/torch/nn/functional.html#pixel_shuffle
|
|
|
"""
|
|
|
def __init__(self, upscale_factor=0.5):
|
|
|
super(PixelSort, self).__init__()
|
|
|
self.upscale_factor = upscale_factor
|
|
|
|
|
|
def forward(self, x):
|
|
|
b, c, h, w = x.size()
|
|
|
x = x.view(b, c, 2, 2, h // 2, w // 2)
|
|
|
x = x.permute(0, 1, 5, 3, 2, 4).contiguous()
|
|
|
x = x.view(b, 4 * c, h // 2, w // 2)
|
|
|
|
|
|
return x
|
|
|
|
|
|
class Downsampler(nn.Sequential):
|
|
|
def __init__(
|
|
|
self, scale, n_feats, bias=True,
|
|
|
conv=default_conv, norm=False, act=False):
|
|
|
|
|
|
modules = []
|
|
|
if scale == 0.5:
|
|
|
modules.append(PixelSort())
|
|
|
modules.append(conv(4 * n_feats, n_feats, 3, bias))
|
|
|
if norm: modules.append(norm(n_feats))
|
|
|
if act: modules.append(act())
|
|
|
else:
|
|
|
raise NotImplementedError
|
|
|
|
|
|
super(Downsampler, self).__init__(*modules)
|
|
|
|
|
|
|
