-
Notifications
You must be signed in to change notification settings - Fork 0
/
model.py
63 lines (46 loc) · 2.21 KB
/
model.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
import torch
import torch.nn as nn
import torchvision.transforms.functional as TF
#Differences from original Unet:
# 1) Add padding to preserve original input size
# 2) Add BatchNorm2d to improve my results
class UNET(nn.Module):
def __init__(self, in_channels, classes):
super(UNET, self).__init__()
self.layers = [in_channels, 64, 128, 256, 512, 1024]
self.double_conv_downs = nn.ModuleList(
[self.__double_conv(layer, layer_n) for layer, layer_n in zip(self.layers[:-1], self.layers[1:])])
self.up_trans = nn.ModuleList(
[nn.ConvTranspose2d(layer, layer_n, kernel_size=2, stride=2)
for layer, layer_n in zip(self.layers[::-1][:-2], self.layers[::-1][1:-1])])
self.double_conv_ups = nn.ModuleList(
[self.__double_conv(layer, layer//2) for layer in self.layers[::-1][:-2]])
self.max_pool_2x2 = nn.MaxPool2d(kernel_size=2, stride=2)
self.final_conv = nn.Conv2d(64, classes, kernel_size=1)
def __double_conv(self, in_channels, out_channels):
conv = nn.Sequential(
nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1, bias=False),
nn.BatchNorm2d(out_channels),
nn.ReLU(inplace=True),
nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
nn.ReLU(inplace=True)
)
return conv
def forward(self, x):
# down layers
concat_layers = []
for down in self.double_conv_downs:
x = down(x)
if down != self.double_conv_downs[-1]:
concat_layers.append(x)
x = self.max_pool_2x2(x)
concat_layers = concat_layers[::-1]
# up layers
for up_trans, double_conv_up, concat_layer in zip(self.up_trans, self.double_conv_ups, concat_layers):
x = up_trans(x)
if x.shape != concat_layer.shape:
x = TF.resize(x, concat_layer.shape[2:])
concatenated = torch.cat((concat_layer, x), dim=1)
x = double_conv_up(concatenated)
x = self.final_conv(x)
return x