File size: 1,027 Bytes
5b9bb29 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 |
import torch.nn as nn
class SubpixelUpsampling(nn.Module):
"""Subpixel Upsampling Module using PixelShuffle"""
def __init__(self, in_channels, scale_factor=2):
super(SubpixelUpsampling, self).__init__()
self.scale_factor = scale_factor
self.conv = nn.Conv2d(
in_channels,
in_channels * (scale_factor ** 2),
kernel_size=3,
padding=1
)
self.pixel_shuffle = nn.PixelShuffle(scale_factor)
def forward(self, x):
x = self.conv(x)
x = self.pixel_shuffle(x)
return x
class UpsamplingBlock(nn.Module):
"""Block for 4x upsampling using two SubpixelUpsampling modules"""
def __init__(self, in_channels):
super(UpsamplingBlock, self).__init__()
self.upsample1 = SubpixelUpsampling(in_channels)
self.upsample2 = SubpixelUpsampling(in_channels)
def forward(self, x):
x = self.upsample1(x)
x = self.upsample2(x)
return x |