| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
|
|
| class FSRCNN(nn.Module): |
| def __init__(self, scale, num_channels=3, d=56, s=12, m=4): |
| super().__init__() |
| self.feature_extraction = nn.Sequential( |
| nn.Conv2d(num_channels, d, kernel_size=5, padding=2), nn.PReLU(d)) |
| self.shrink = nn.Sequential(nn.Conv2d(d, s, kernel_size=1), nn.PReLU(s)) |
| mapping_layers = [] |
| for _ in range(m): |
| mapping_layers += [nn.Conv2d(s, s, kernel_size=3, padding=1), nn.PReLU(s)] |
| self.mapping = nn.Sequential(*mapping_layers) |
| self.expand = nn.Sequential(nn.Conv2d(s, d, kernel_size=1), nn.PReLU(d)) |
| self.upsample = nn.Sequential( |
| nn.Conv2d(d, num_channels * scale * scale, kernel_size=3, padding=1), |
| nn.PixelShuffle(scale)) |
| self.scale = scale |
|
|
| def forward(self, x): |
| x = self.feature_extraction(x) |
| x = self.shrink(x) |
| x = self.mapping(x) |
| x = self.expand(x) |
| x = self.upsample(x) |
| return torch.sigmoid(x) |
|
|
