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