-
Notifications
You must be signed in to change notification settings - Fork 16
/
FCN.py
132 lines (117 loc) · 4.37 KB
/
FCN.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
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
# -*- coding: utf-8 -*-
import torch
import torch.nn as nn
class Imagenet_Encoder(nn.Module):
def __init__(self):
super().__init__()
self.conv1_1 = nn.Sequential(
nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, stride=1, padding=1),
nn.ReLU(inplace=True),
nn.BatchNorm2d(16)
)
self.conv1_2 = nn.Sequential(
nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, stride=1, padding=1),
nn.ReLU(inplace=True),
nn.BatchNorm2d(32),
nn.MaxPool2d(kernel_size=2)
)
self.conv2_1 = nn.Sequential(
nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, stride=1, padding=1),
nn.ReLU(inplace=True),
nn.BatchNorm2d(64)
)
self.conv2_2 = nn.Sequential(
nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1),
nn.ReLU(inplace=True),
nn.BatchNorm2d(64),
nn.MaxPool2d(kernel_size=2)
)
self.conv3_1 = nn.Sequential(
nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, stride=1, padding=1),
nn.ReLU(inplace=True),
nn.BatchNorm2d(128)
)
self.conv3_2 = nn.Sequential(
nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, stride=1, padding=1),
nn.ReLU(inplace=True),
nn.BatchNorm2d(128),
nn.MaxPool2d(kernel_size=2)
)
self.conv4_1 = nn.Sequential(
nn.Conv2d(in_channels=128, out_channels=32, kernel_size=3, stride=1, padding=1),
nn.ReLU(inplace=True),
nn.BatchNorm2d(32)
)
self.conv4_2 = nn.Sequential(
nn.Conv2d(in_channels=32, out_channels=8, kernel_size=3, stride=1, padding=1),
nn.ReLU(inplace=True),
nn.BatchNorm2d(8),
nn.MaxPool2d(kernel_size=2)
)
def forward(self, x):
x = self.conv1_1(x)
x = self.conv1_2(x)
x = self.conv2_1(x)
x = self.conv2_2(x)
x = self.conv3_1(x)
x = self.conv3_2(x)
x = self.conv4_1(x)
x = self.conv4_2(x)
return x
class Imagenet_Decoder(nn.Module):
def __init__(self):
super().__init__()
self.deconv1_1 = nn.Sequential(
nn.Conv2d(in_channels=8, out_channels=32, kernel_size=3, stride=1, padding=1),
nn.ReLU(inplace=True),
nn.BatchNorm2d(32)
)
self.deconv1_2 = nn.Sequential(
nn.ConvTranspose2d(32, 64, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1),
nn.ReLU(inplace=True),
nn.BatchNorm2d(64)
)
self.deconv2_1 = nn.Sequential(
nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, stride=1, padding=1),
nn.ReLU(inplace=True),
nn.BatchNorm2d(128)
)
self.deconv2_2 = nn.Sequential(
nn.ConvTranspose2d(128, 128, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1),
nn.ReLU(inplace=True),
nn.BatchNorm2d(128)
)
self.deconv3_1 = nn.Sequential(
nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, stride=1, padding=1),
nn.ReLU(inplace=True),
nn.BatchNorm2d(128)
)
self.deconv3_2 = nn.Sequential(
nn.ConvTranspose2d(128, 64, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1),
nn.ReLU(inplace=True),
nn.BatchNorm2d(64)
)
self.deconv4_1 = nn.Sequential(
nn.Conv2d(in_channels=64, out_channels=32, kernel_size=3, stride=1, padding=1),
nn.ReLU(inplace=True),
nn.BatchNorm2d(32)
)
self.deconv4_2 = nn.Sequential(
nn.ConvTranspose2d(32, 16, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1),
nn.ReLU(inplace=True),
nn.BatchNorm2d(16)
)
self.fcn = nn.Conv2d(16, 3, kernel_size=1)
def forward(self, x):
n = len(x)
x = x.view(n, 8, 14, 14)
y = self.deconv1_1(x)
y = self.deconv1_2(y)
y = self.deconv2_1(y)
y = self.deconv2_2(y)
y = self.deconv3_1(y)
y = self.deconv3_2(y)
y = self.deconv4_1(y)
y = self.deconv4_2(y)
y = self.fcn(y)
return torch.tanh(y)