Spaces:
Running
Running
| import math | |
| import torch | |
| from torch import nn as nn | |
| from torch.nn import functional as F | |
| from torch.nn import init as init | |
| from torch.nn.modules.batchnorm import _BatchNorm | |
| def default_init_weights(module_list, scale=1, bias_fill=0, **kwargs): | |
| """Initialize network weights. | |
| Args: | |
| module_list (list[nn.Module] | nn.Module): Modules to be initialized. | |
| scale (float): Scale initialized weights, especially for residual | |
| blocks. Default: 1. | |
| bias_fill (float): The value to fill bias. Default: 0 | |
| kwargs (dict): Other arguments for initialization function. | |
| """ | |
| if not isinstance(module_list, list): | |
| module_list = [module_list] | |
| for module in module_list: | |
| for m in module.modules(): | |
| if isinstance(m, nn.Conv2d): | |
| init.kaiming_normal_(m.weight, **kwargs) | |
| m.weight.data *= scale | |
| if m.bias is not None: | |
| m.bias.data.fill_(bias_fill) | |
| elif isinstance(m, nn.Linear): | |
| init.kaiming_normal_(m.weight, **kwargs) | |
| m.weight.data *= scale | |
| if m.bias is not None: | |
| m.bias.data.fill_(bias_fill) | |
| elif isinstance(m, _BatchNorm): | |
| init.constant_(m.weight, 1) | |
| if m.bias is not None: | |
| m.bias.data.fill_(bias_fill) | |
| def make_layer(basic_block, num_basic_block, **kwarg): | |
| """Make layers by stacking the same blocks. | |
| Args: | |
| basic_block (nn.module): nn.module class for basic block. | |
| num_basic_block (int): number of blocks. | |
| Returns: | |
| nn.Sequential: Stacked blocks in nn.Sequential. | |
| """ | |
| layers = [] | |
| for _ in range(num_basic_block): | |
| layers.append(basic_block(**kwarg)) | |
| return nn.Sequential(*layers) | |
| class ResidualBlockNoBN(nn.Module): | |
| """Residual block without BN. | |
| It has a style of: | |
| ---Conv-ReLU-Conv-+- | |
| |________________| | |
| Args: | |
| num_feat (int): Channel number of intermediate features. | |
| Default: 64. | |
| res_scale (float): Residual scale. Default: 1. | |
| pytorch_init (bool): If set to True, use pytorch default init, | |
| otherwise, use default_init_weights. Default: False. | |
| """ | |
| def __init__(self, num_feat=64, res_scale=1, pytorch_init=False): | |
| super(ResidualBlockNoBN, self).__init__() | |
| self.res_scale = res_scale | |
| self.conv1 = nn.Conv2d(num_feat, num_feat, 3, 1, 1, bias=True) | |
| self.conv2 = nn.Conv2d(num_feat, num_feat, 3, 1, 1, bias=True) | |
| self.relu = nn.ReLU(inplace=True) | |
| if not pytorch_init: | |
| default_init_weights([self.conv1, self.conv2], 0.1) | |
| def forward(self, x): | |
| identity = x | |
| out = self.conv2(self.relu(self.conv1(x))) | |
| return identity + out * self.res_scale | |
| class Upsample(nn.Sequential): | |
| """Upsample module. | |
| Args: | |
| scale (int): Scale factor. Supported scales: 2^n and 3. | |
| num_feat (int): Channel number of intermediate features. | |
| """ | |
| def __init__(self, scale, num_feat): | |
| m = [] | |
| if (scale & (scale - 1)) == 0: # scale = 2^n | |
| for _ in range(int(math.log(scale, 2))): | |
| m.append(nn.Conv2d(num_feat, 4 * num_feat, 3, 1, 1)) | |
| m.append(nn.PixelShuffle(2)) | |
| elif scale == 3: | |
| m.append(nn.Conv2d(num_feat, 9 * num_feat, 3, 1, 1)) | |
| m.append(nn.PixelShuffle(3)) | |
| else: | |
| raise ValueError(f'scale {scale} is not supported. ' | |
| 'Supported scales: 2^n and 3.') | |
| super(Upsample, self).__init__(*m) | |
| # TODO: may write a cpp file | |
| def pixel_unshuffle(x, scale): | |
| """ Pixel unshuffle. | |
| Args: | |
| x (Tensor): Input feature with shape (b, c, hh, hw). | |
| scale (int): Downsample ratio. | |
| Returns: | |
| Tensor: the pixel unshuffled feature. | |
| """ | |
| b, c, hh, hw = x.size() | |
| out_channel = c * (scale**2) | |
| assert hh % scale == 0 and hw % scale == 0 | |
| h = hh // scale | |
| w = hw // scale | |
| x_view = x.view(b, c, h, scale, w, scale) | |
| return x_view.permute(0, 1, 3, 5, 2, 4).reshape(b, out_channel, h, w) |