261 lines
11 KiB
Python
261 lines
11 KiB
Python
import torch
|
||
import torch.nn as nn
|
||
from torchvision import models
|
||
from torchvision.models.vgg import VGG
|
||
|
||
"""Architecture选择：FCN32s; FCN16s; FCN8s; FCNs"""
|
||
class FCN32s(nn.Module):
|
||
|
||
def __init__(self, pretrained_net, n_class):
|
||
super().__init__()
|
||
self.n_class = n_class
|
||
self.pretrained_net = pretrained_net
|
||
self.relu = nn.ReLU(inplace=True)
|
||
self.deconv1 = nn.ConvTranspose2d(512, 512, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1)
|
||
self.bn1 = nn.BatchNorm2d(512)
|
||
self.deconv2 = nn.ConvTranspose2d(512, 256, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1)
|
||
self.bn2 = nn.BatchNorm2d(256)
|
||
self.deconv3 = nn.ConvTranspose2d(256, 128, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1)
|
||
self.bn3 = nn.BatchNorm2d(128)
|
||
self.deconv4 = nn.ConvTranspose2d(128, 64, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1)
|
||
self.bn4 = nn.BatchNorm2d(64)
|
||
self.deconv5 = nn.ConvTranspose2d(64, 32, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1)
|
||
self.bn5 = nn.BatchNorm2d(32)
|
||
self.classifier = nn.Conv2d(32, n_class, kernel_size=1)
|
||
|
||
def forward(self, x):
|
||
output = self.pretrained_net(x)
|
||
x5 = output['x5']
|
||
|
||
score = self.bn1(self.relu(self.deconv1(x5)))
|
||
score = self.bn2(self.relu(self.deconv2(score)))
|
||
score = self.bn3(self.relu(self.deconv3(score)))
|
||
score = self.bn4(self.relu(self.deconv4(score)))
|
||
score = self.bn5(self.relu(self.deconv5(score)))
|
||
score = self.classifier(score)
|
||
|
||
return score
|
||
|
||
|
||
class FCN16s(nn.Module):
|
||
|
||
def __init__(self, pretrained_net, n_class):
|
||
super().__init__()
|
||
self.n_class = n_class
|
||
self.pretrained_net = pretrained_net
|
||
self.relu = nn.ReLU(inplace=True)
|
||
self.deconv1 = nn.ConvTranspose2d(512, 512, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1)
|
||
self.bn1 = nn.BatchNorm2d(512)
|
||
self.deconv2 = nn.ConvTranspose2d(512, 256, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1)
|
||
self.bn2 = nn.BatchNorm2d(256)
|
||
self.deconv3 = nn.ConvTranspose2d(256, 128, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1)
|
||
self.bn3 = nn.BatchNorm2d(128)
|
||
self.deconv4 = nn.ConvTranspose2d(128, 64, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1)
|
||
self.bn4 = nn.BatchNorm2d(64)
|
||
self.deconv5 = nn.ConvTranspose2d(64, 32, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1)
|
||
self.bn5 = nn.BatchNorm2d(32)
|
||
self.classifier = nn.Conv2d(32, n_class, kernel_size=1)
|
||
|
||
def forward(self, x):
|
||
output = self.pretrained_net(x)
|
||
x5 = output['x5']
|
||
x4 = output['x4']
|
||
|
||
score = self.relu(self.deconv1(x5))
|
||
score = self.bn1(score + x4)
|
||
score = self.bn2(self.relu(self.deconv2(score)))
|
||
score = self.bn3(self.relu(self.deconv3(score)))
|
||
score = self.bn4(self.relu(self.deconv4(score)))
|
||
score = self.bn5(self.relu(self.deconv5(score)))
|
||
score = self.classifier(score)
|
||
|
||
return score
|
||
|
||
|
||
class FCN8s(nn.Module):
|
||
|
||
def __init__(self, pretrained_net, n_class):
|
||
super().__init__()
|
||
self.n_class = n_class
|
||
self.pretrained_net = pretrained_net
|
||
self.relu = nn.ReLU(inplace=True)
|
||
self.deconv1 = nn.ConvTranspose2d(512, 512, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1)
|
||
self.bn1 = nn.BatchNorm2d(512)
|
||
self.deconv2 = nn.ConvTranspose2d(512, 256, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1)
|
||
self.bn2 = nn.BatchNorm2d(256)
|
||
self.deconv3 = nn.ConvTranspose2d(256, 128, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1)
|
||
self.bn3 = nn.BatchNorm2d(128)
|
||
self.deconv4 = nn.ConvTranspose2d(128, 64, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1)
|
||
self.bn4 = nn.BatchNorm2d(64)
|
||
self.deconv5 = nn.ConvTranspose2d(64, 32, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1)
|
||
self.bn5 = nn.BatchNorm2d(32)
|
||
self.classifier = nn.Conv2d(32, n_class, kernel_size=1)
|
||
|
||
def forward(self, x):
|
||
output = self.pretrained_net(x)
|
||
x5 = output['x5']
|
||
x4 = output['x4']
|
||
x3 = output['x3']
|
||
|
||
score = self.relu(self.deconv1(x5))
|
||
score = self.bn1(score + x4)
|
||
score = self.relu(self.deconv2(score))
|
||
score = self.bn2(score + x3)
|
||
score = self.bn3(self.relu(self.deconv3(score)))
|
||
score = self.bn4(self.relu(self.deconv4(score)))
|
||
score = self.bn5(self.relu(self.deconv5(score)))
|
||
score = self.classifier(score)
|
||
|
||
return score
|
||
|
||
|
||
class FCNs(nn.Module):
|
||
|
||
def __init__(self, pretrained_net, n_class):
|
||
super().__init__()
|
||
self.n_class = n_class
|
||
self.pretrained_net = pretrained_net
|
||
self.relu = nn.ReLU(inplace=True)
|
||
self.deconv1 = nn.ConvTranspose2d(512, 512, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1)
|
||
self.bn1 = nn.BatchNorm2d(512)
|
||
self.deconv2 = nn.ConvTranspose2d(512, 256, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1)
|
||
self.bn2 = nn.BatchNorm2d(256)
|
||
self.deconv3 = nn.ConvTranspose2d(256, 128, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1)
|
||
self.bn3 = nn.BatchNorm2d(128)
|
||
self.deconv4 = nn.ConvTranspose2d(128, 64, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1)
|
||
self.bn4 = nn.BatchNorm2d(64)
|
||
self.deconv5 = nn.ConvTranspose2d(64, 32, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1)
|
||
self.bn5 = nn.BatchNorm2d(32)
|
||
self.classifier = nn.Conv2d(32, n_class, kernel_size=1)
|
||
# classifier is 1x1 conv, to reduce channels from 32 to n_class
|
||
|
||
def forward(self, x):
|
||
output = self.pretrained_net(x)
|
||
x5 = output['x5']
|
||
x4 = output['x4']
|
||
x3 = output['x3']
|
||
x2 = output['x2']
|
||
x1 = output['x1']
|
||
|
||
score = self.bn1(self.relu(self.deconv1(x5)))
|
||
score = score + x4
|
||
score = self.bn2(self.relu(self.deconv2(score)))
|
||
score = score + x3
|
||
score = self.bn3(self.relu(self.deconv3(score)))
|
||
score = score + x2
|
||
score = self.bn4(self.relu(self.deconv4(score)))
|
||
score = score + x1
|
||
score = self.bn5(self.relu(self.deconv5(score)))
|
||
score = self.classifier(score)
|
||
|
||
return score
|
||
|
||
"""pretrained:是否使用预训练模型；model:backbone类型（vgg11;vgg13;vgg16;vgg19）；其他参数使用默认设置"""
|
||
class VGGNet(VGG):
|
||
def __init__(self, pretrained=True, model='vgg16', requires_grad=True, remove_fc=True, show_params=False):
|
||
super().__init__(make_layers(cfg[model]))
|
||
self.ranges = ranges[model]
|
||
|
||
if pretrained:
|
||
exec("self.load_state_dict(models.%s(pretrained=True).state_dict())" % model)
|
||
|
||
if not requires_grad:
|
||
for param in super().parameters():
|
||
param.requires_grad = False
|
||
|
||
# delete redundant fully-connected layer params, can save memory
|
||
# 去掉vgg最后的全连接层(classifier)
|
||
if remove_fc:
|
||
del self.classifier
|
||
|
||
if show_params:
|
||
for name, param in self.named_parameters():
|
||
print(name, param.size())
|
||
|
||
def forward(self, x):
|
||
output = {}
|
||
# get the output of each maxpooling layer (5 maxpool in VGG net)
|
||
for idx, (begin, end) in enumerate(self.ranges):
|
||
#self.ranges = ((0, 5), (5, 10), (10, 17), (17, 24), (24, 31)) (vgg16 examples)
|
||
for layer in range(begin, end):
|
||
x = self.features[layer](x)
|
||
output["x%d"%(idx+1)] = x
|
||
|
||
return output
|
||
|
||
|
||
ranges = {
|
||
'vgg11': ((0, 3), (3, 6), (6, 11), (11, 16), (16, 21)),
|
||
'vgg13': ((0, 5), (5, 10), (10, 15), (15, 20), (20, 25)),
|
||
'vgg16': ((0, 5), (5, 10), (10, 17), (17, 24), (24, 31)),
|
||
'vgg19': ((0, 5), (5, 10), (10, 19), (19, 28), (28, 37))
|
||
}
|
||
|
||
# Vgg-Net config
|
||
# Vgg网络结构配置
|
||
cfg = {
|
||
'vgg11': [64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
|
||
'vgg13': [64, 64, 'M', 128, 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
|
||
'vgg16': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M'],
|
||
'vgg19': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512, 'M'],
|
||
}
|
||
|
||
# make layers using Vgg-Net config(cfg)
|
||
# 由cfg构建vgg-Net
|
||
def make_layers(cfg, batch_norm=False):
|
||
layers = []
|
||
in_channels = 3
|
||
for v in cfg:
|
||
if v == 'M':
|
||
layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
|
||
else:
|
||
conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1)
|
||
if batch_norm:
|
||
layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)]
|
||
else:
|
||
layers += [conv2d, nn.ReLU(inplace=True)]
|
||
in_channels = v
|
||
return nn.Sequential(*layers)
|
||
|
||
'''
|
||
VGG-16网络参数
|
||
|
||
Sequential(
|
||
(0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
|
||
(1): ReLU(inplace)
|
||
(2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
|
||
(3): ReLU(inplace)
|
||
(4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
|
||
(5): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
|
||
(6): ReLU(inplace)
|
||
(7): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
|
||
(8): ReLU(inplace)
|
||
(9): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
|
||
(10): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
|
||
(11): ReLU(inplace)
|
||
(12): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
|
||
(13): ReLU(inplace)
|
||
(14): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
|
||
(15): ReLU(inplace)
|
||
(16): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
|
||
(17): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
|
||
(18): ReLU(inplace)
|
||
(19): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
|
||
(20): ReLU(inplace)
|
||
(21): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
|
||
(22): ReLU(inplace)
|
||
(23): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
|
||
(24): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
|
||
(25): ReLU(inplace)
|
||
(26): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
|
||
(27): ReLU(inplace)
|
||
(28): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
|
||
(29): ReLU(inplace)
|
||
(30): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
|
||
)
|
||
|
||
'''
|
||
|
||
|
||
if __name__ == "__main__":
|
||
pass |