pytorch实现简单的FCN全卷积网络结构，用于语义分割

import torch import torch.nn as nn from torchvision import models from torchvision.models.vgg import VGG 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 class VGGNet(VGG): def __init__(self, pretrained=False, 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,