|
|
import collections.abc |
|
|
import math |
|
|
import torch |
|
|
import torchvision |
|
|
import warnings |
|
|
from itertools import repeat |
|
|
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 |
|
|
|
|
|
@torch.no_grad() |
|
|
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. |
|
|
|
|
|
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 EDSRNOUP(nn.Module): |
|
|
def __init__(self, |
|
|
num_in_ch=3, |
|
|
num_out_ch=3, |
|
|
num_feat=64, |
|
|
num_block=16, |
|
|
upscale=4, |
|
|
res_scale=1): |
|
|
super(EDSRNOUP, self).__init__() |
|
|
|
|
|
self.conv_first = nn.Conv2d(num_in_ch, num_feat, 3, 1, 1) |
|
|
self.body = make_layer(ResidualBlockNoBN, num_block, num_feat=num_feat, res_scale=res_scale, pytorch_init=True) |
|
|
self.conv_after_body = nn.Conv2d(num_feat, num_feat, 3, 1, 1) |
|
|
|
|
|
|
|
|
def forward(self, x): |
|
|
|
|
|
x = self.conv_first(x) |
|
|
res = self.conv_after_body(self.body(x)) |
|
|
x = res + x |
|
|
|
|
|
return res |
|
|
|
|
|
|
|
|
if __name__ == '__main__': |
|
|
x = torch.randn(8,3,48,48) |
|
|
model = EDSRNOUP(num_in_ch=3, num_out_ch=3) |
|
|
y = model(x) |
|
|
print(y.shape) |
