| from model import common
|
|
|
| import torch.nn as nn
|
| import torch.nn.init as init
|
|
|
| url = {
|
| 'r20f64': ''
|
| }
|
|
|
| def make_model(args, parent=False):
|
| return VDSR(args)
|
|
|
| class VDSR(nn.Module):
|
| def __init__(self, args, conv=common.default_conv):
|
| super(VDSR, self).__init__()
|
|
|
| n_resblocks = args.n_resblocks
|
| n_feats = args.n_feats
|
| kernel_size = 3
|
| self.url = url['r{}f{}'.format(n_resblocks, n_feats)]
|
| self.sub_mean = common.MeanShift(args.rgb_range)
|
| self.add_mean = common.MeanShift(args.rgb_range, sign=1)
|
|
|
| def basic_block(in_channels, out_channels, act):
|
| return common.BasicBlock(
|
| conv, in_channels, out_channels, kernel_size,
|
| bias=True, bn=False, act=act
|
| )
|
|
|
|
|
| m_body = []
|
| m_body.append(basic_block(args.n_colors, n_feats, nn.ReLU(True)))
|
| for _ in range(n_resblocks - 2):
|
| m_body.append(basic_block(n_feats, n_feats, nn.ReLU(True)))
|
| m_body.append(basic_block(n_feats, args.n_colors, None))
|
|
|
| self.body = nn.Sequential(*m_body)
|
|
|
| def forward(self, x):
|
| x = self.sub_mean(x)
|
| res = self.body(x)
|
| res += x
|
| x = self.add_mean(res)
|
|
|
| return x
|
|
|
|
|
