FCN模型讲解

1
import torch
2
import torch.nn as nn
3
import torch.nn.functional as F
4
import numpy as np

1
class FCN_VGG16(nn.Module):
2
    '''
3
    FCN 的 backbone，由 VGG16 修改而来，舍弃最后的全连接层
4
    以池化层为区分，一个池化层到上一个池化层之间的部分认为一个卷积块。
5
    '''
6
    def __init__(self):
7
        super(FCN_VGG16, self).__init__()
8
        self.features = nn.Sequential(
9
            # 第一个卷积块: 输入通道数：3，输出通道数：64，卷积核大小：3*3，步长：1，填充：1
10
            nn.Conv2d(in_channels=3, out_channels=64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
11
            nn.ReLU(inplace=True),
12
            nn.Conv2d(in_channels=64, out_channels=64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
13
            nn.ReLU(inplace=True),
14
            nn.MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False),
15
            # 第二个卷积块
16
            nn.Conv2d(in_channels=64, out_channels=128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
17
            nn.ReLU(inplace=True),
18
            nn.Conv2d(in_channels=128, out_channels=128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
19
            nn.ReLU(inplace=True),
20
            nn.MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False),
21
            # 第三个卷积块
22
            nn.Conv2d(in_channels=128, out_channels=256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
23
            nn.ReLU(inplace=True),
24
            nn.Conv2d(in_channels=256, out_channels=256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
25
            nn.ReLU(inplace=True),
26
            nn.Conv2d(in_channels=256, out_channels=256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
27
            nn.ReLU(inplace=True),
28
            nn.MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False),
29
            # 第四个卷积块
30
            nn.Conv2d(in_channels=256, out_channels=512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
31
            nn.ReLU(inplace=True),
32
            nn.Conv2d(in_channels=512,out_channels=512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
33
            nn.ReLU(inplace=True),
34
            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
35
            nn.ReLU(inplace=True),
36
            nn.MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False),
37
            # 第五个卷积块
38
            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
39
            nn.ReLU(inplace=True),
40
            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
41
            nn.ReLU(inplace=True),
42
            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
43
            nn.ReLU(inplace=True),
44
            nn.MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False),
45
        )
46

47
        # 每一层在 features 中的范围，{0，1，2，3，4} 为第一个卷积块，{5，6，7，8，9} 为第二个卷积块...
48
        self.range = ((0, 5), (5, 10), (10, 17), (17, 24), (24, 31))
49

50
    def forward(self, input):
51
        output = {}
52
        # 每一块的输出
53
        for idx, (start, end) in enumerate(self.range):
54
            for layer in range(start, end):
55
                input = self.features[layer](input)
56
            output["x%d" % (idx + 1)] = input
57
        return output

1
def test_vgg():
2
    # Backbone 的测试函数
3
    input_x = torch.randn((1,3,512,512))
4
    vgg = FCN_VGG16()
5
    output_y = vgg(input_x)
6

7
    for key in output_y:
8
        print(output_y[key].size())
9

10
test_vgg()