| |
| import torch.nn as nn |
| from math import sqrt |
|
|
| class Conv_ReLU_Block(nn.Module): |
| def __init__(self): |
| super(Conv_ReLU_Block, self).__init__() |
| self.conv = nn.Conv2d(64, 64, kernel_size=3, padding=1, bias=False) |
| self.relu = nn.ReLU(inplace=True) |
| |
| def forward(self, x): |
| return self.relu(self.conv(x)) |
|
|
| class VDSR(nn.Module): |
| def __init__(self): |
| super(VDSR, self).__init__() |
| self.residual_layer = self.make_layer(Conv_ReLU_Block, 18) |
| self.input = nn.Conv2d(1, 64, kernel_size=3, padding=1, bias=False) |
| self.output = nn.Conv2d(64, 1, kernel_size=3, padding=1, bias=False) |
| self.relu = nn.ReLU(inplace=True) |
| |
| |
| for m in self.modules(): |
| if isinstance(m, nn.Conv2d): |
| n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels |
| m.weight.data.normal_(0, sqrt(2. / n)) |
| |
| def make_layer(self, block, num_of_layer): |
| layers = [] |
| for _ in range(num_of_layer): |
| layers.append(block()) |
| return nn.Sequential(*layers) |
|
|
| def forward(self, x): |
| residual = x |
| out = self.relu(self.input(x)) |
| out = self.residual_layer(out) |
| out = self.output(out) |
| return out + residual |
