-
Notifications
You must be signed in to change notification settings - Fork 392
/
Copy pathseg_net.py
69 lines (60 loc) · 2.51 KB
/
seg_net.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
import torch
from torch import nn
from torchvision import models
from ..utils import initialize_weights
class _DecoderBlock(nn.Module):
def __init__(self, in_channels, out_channels, num_conv_layers):
super(_DecoderBlock, self).__init__()
middle_channels = in_channels // 2
layers = [
nn.ConvTranspose2d(in_channels, in_channels, kernel_size=2, stride=2),
nn.Conv2d(in_channels, middle_channels, kernel_size=3, padding=1),
nn.BatchNorm2d(middle_channels),
nn.ReLU(inplace=True)
]
layers += [
nn.Conv2d(middle_channels, middle_channels, kernel_size=3, padding=1),
nn.BatchNorm2d(middle_channels),
nn.ReLU(inplace=True),
] * (num_conv_layers - 2)
layers += [
nn.Conv2d(middle_channels, out_channels, kernel_size=3, padding=1),
nn.BatchNorm2d(out_channels),
nn.ReLU(inplace=True),
]
self.decode = nn.Sequential(*layers)
def forward(self, x):
return self.decode(x)
class SegNet(nn.Module):
def __init__(self, num_classes, pretrained=True):
super(SegNet, self).__init__()
vgg = models.vgg19_bn(pretrained=pretrained)
features = list(vgg.features.children())
self.enc1 = nn.Sequential(*features[0:7])
self.enc2 = nn.Sequential(*features[7:14])
self.enc3 = nn.Sequential(*features[14:27])
self.enc4 = nn.Sequential(*features[27:40])
self.enc5 = nn.Sequential(*features[40:])
self.dec5 = nn.Sequential(
*([nn.ConvTranspose2d(512, 512, kernel_size=2, stride=2)] +
[nn.Conv2d(512, 512, kernel_size=3, padding=1),
nn.BatchNorm2d(512),
nn.ReLU(inplace=True)] * 4)
)
self.dec4 = _DecoderBlock(1024, 256, 4)
self.dec3 = _DecoderBlock(512, 128, 4)
self.dec2 = _DecoderBlock(256, 64, 2)
self.dec1 = _DecoderBlock(128, num_classes, 2)
initialize_weights(self.dec5, self.dec4, self.dec3, self.dec2, self.dec1)
def forward(self, x):
enc1 = self.enc1(x)
enc2 = self.enc2(enc1)
enc3 = self.enc3(enc2)
enc4 = self.enc4(enc3)
enc5 = self.enc5(enc4)
dec5 = self.dec5(enc5)
dec4 = self.dec4(torch.cat([enc4, dec5], 1))
dec3 = self.dec3(torch.cat([enc3, dec4], 1))
dec2 = self.dec2(torch.cat([enc2, dec3], 1))
dec1 = self.dec1(torch.cat([enc1, dec2], 1))
return dec1