-
Notifications
You must be signed in to change notification settings - Fork 92
/
models.py
116 lines (87 loc) · 3.81 KB
/
models.py
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
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
# -*- coding: utf-8 -*-
"""Implements SRGAN models: https://arxiv.org/abs/1609.04802
TODO:
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
def swish(x):
return x * F.sigmoid(x)
class FeatureExtractor(nn.Module):
def __init__(self, cnn, feature_layer=11):
super(FeatureExtractor, self).__init__()
self.features = nn.Sequential(*list(cnn.features.children())[:(feature_layer+1)])
def forward(self, x):
return self.features(x)
class residualBlock(nn.Module):
def __init__(self, in_channels=64, k=3, n=64, s=1):
super(residualBlock, self).__init__()
self.conv1 = nn.Conv2d(in_channels, n, k, stride=s, padding=1)
self.bn1 = nn.BatchNorm2d(n)
self.conv2 = nn.Conv2d(n, n, k, stride=s, padding=1)
self.bn2 = nn.BatchNorm2d(n)
def forward(self, x):
y = swish(self.bn1(self.conv1(x)))
return self.bn2(self.conv2(y)) + x
class upsampleBlock(nn.Module):
# Implements resize-convolution
def __init__(self, in_channels, out_channels):
super(upsampleBlock, self).__init__()
self.conv = nn.Conv2d(in_channels, out_channels, 3, stride=1, padding=1)
self.shuffler = nn.PixelShuffle(2)
def forward(self, x):
return swish(self.shuffler(self.conv(x)))
class Generator(nn.Module):
def __init__(self, n_residual_blocks, upsample_factor):
super(Generator, self).__init__()
self.n_residual_blocks = n_residual_blocks
self.upsample_factor = upsample_factor
self.conv1 = nn.Conv2d(3, 64, 9, stride=1, padding=4)
for i in range(self.n_residual_blocks):
self.add_module('residual_block' + str(i+1), residualBlock())
self.conv2 = nn.Conv2d(64, 64, 3, stride=1, padding=1)
self.bn2 = nn.BatchNorm2d(64)
for i in range(self.upsample_factor/2):
self.add_module('upsample' + str(i+1), upsampleBlock(64, 256))
self.conv3 = nn.Conv2d(64, 3, 9, stride=1, padding=4)
def forward(self, x):
x = swish(self.conv1(x))
y = x.clone()
for i in range(self.n_residual_blocks):
y = self.__getattr__('residual_block' + str(i+1))(y)
x = self.bn2(self.conv2(y)) + x
for i in range(self.upsample_factor/2):
x = self.__getattr__('upsample' + str(i+1))(x)
return self.conv3(x)
class Discriminator(nn.Module):
def __init__(self):
super(Discriminator, self).__init__()
self.conv1 = nn.Conv2d(3, 64, 3, stride=1, padding=1)
self.conv2 = nn.Conv2d(64, 64, 3, stride=2, padding=1)
self.bn2 = nn.BatchNorm2d(64)
self.conv3 = nn.Conv2d(64, 128, 3, stride=1, padding=1)
self.bn3 = nn.BatchNorm2d(128)
self.conv4 = nn.Conv2d(128, 128, 3, stride=2, padding=1)
self.bn4 = nn.BatchNorm2d(128)
self.conv5 = nn.Conv2d(128, 256, 3, stride=1, padding=1)
self.bn5 = nn.BatchNorm2d(256)
self.conv6 = nn.Conv2d(256, 256, 3, stride=2, padding=1)
self.bn6 = nn.BatchNorm2d(256)
self.conv7 = nn.Conv2d(256, 512, 3, stride=1, padding=1)
self.bn7 = nn.BatchNorm2d(512)
self.conv8 = nn.Conv2d(512, 512, 3, stride=2, padding=1)
self.bn8 = nn.BatchNorm2d(512)
# Replaced original paper FC layers with FCN
self.conv9 = nn.Conv2d(512, 1, 1, stride=1, padding=1)
def forward(self, x):
x = swish(self.conv1(x))
x = swish(self.bn2(self.conv2(x)))
x = swish(self.bn3(self.conv3(x)))
x = swish(self.bn4(self.conv4(x)))
x = swish(self.bn5(self.conv5(x)))
x = swish(self.bn6(self.conv6(x)))
x = swish(self.bn7(self.conv7(x)))
x = swish(self.bn8(self.conv8(x)))
x = self.conv9(x)
return F.sigmoid(F.avg_pool2d(x, x.size()[2:])).view(x.size()[0], -1)