Spaces:
Runtime error
Runtime error
File size: 2,799 Bytes
7c4166f |
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 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 |
import torch
from torch import nn
class Residual_block(nn.Module):
"""Residual block
Architecture: https://arxiv.org/pdf/1610.02915.pdf
"""
def __init__(self, channel):
super(Residual_block, self).__init__()
self.conv_1 = nn.Conv2d(in_channels=channel, out_channels=channel,
padding='same', kernel_size=3, stride=1)
self.inst1 = nn.InstanceNorm2d(channel, affine=True)
self.conv_2 = nn.Conv2d(in_channels=channel, out_channels=channel,
padding='same', kernel_size=3, stride=1)
self.inst2 = nn.InstanceNorm2d(channel, affine=True)
self.relu = nn.ReLU()
def forward(self, x):
residual = x
out = self.relu(self.inst1(self.conv_1(x)))
out = self.inst2(self.conv_2(out))
return self.relu(out + residual)
class TransformerNet(nn.Module):
def __init__(self):
super(TransformerNet, self).__init__()
# Downsampling
self.conv1 = nn.Conv2d(in_channels=3, out_channels=32, kernel_size=9, stride=1, padding = 9//2)
self.BN_1 = nn.InstanceNorm2d(num_features=32, affine=True)
self.down_1 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, stride=2, padding = 1)
self.BN_2 = nn.InstanceNorm2d(num_features=64, affine=True)
self.down_2 = nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, stride=2, padding = 1)
self.BN_3 = nn.InstanceNorm2d(num_features=128, affine=True)
# Residual connect
self.res_1 = Residual_block(128)
self.res_2 = Residual_block(128)
self.res_3 = Residual_block(128)
self.res_4 = Residual_block(128)
self.res_5 = Residual_block(128)
# Upsampling
self.up_1 = nn.ConvTranspose2d(in_channels=128, out_channels=64, kernel_size=3, stride=2, padding=1, output_padding= 1)
self.BN_4 = nn.InstanceNorm2d(num_features=64, affine=True)
self.up_2 = nn.ConvTranspose2d(in_channels=64, out_channels=32, kernel_size=3, stride=2, padding = 1, output_padding= 1)
self.BN_5 = nn.InstanceNorm2d(num_features=32, affine=True)
self.conv2 = nn.Conv2d(in_channels=32, out_channels=3, kernel_size=9, stride=1, padding = 9//2)
self.relu = nn.ReLU()
def forward(self, x):
y = self.relu(self.BN_1(self.conv1(x)))
# print(y.shape)
y = self.relu(self.BN_2(self.down_1(y)))
# print(y.shape)
y = self.relu(self.BN_3(self.down_2(y)))
# print(y.shape)
# print()
y = self.res_1(y)
# print(y.shape)
y = self.res_2(y)
# print(y.shape)
y = self.res_3(y)
# print(y.shape)
y = self.res_4(y)
# print(y.shape)
y = self.res_5(y)
# print(y.shape)
# print()
y = self.relu(self.BN_4(self.up_1(y)))
# print(y.shape)
y = self.relu(self.BN_5(self.up_2(y)))
# print(y.shape)
y = self.conv2(y)
# print(y.shape)
return y
|