| | 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)
|
| |
|
| |
|
