复现InceptionV1

Going deeper with convolutions
Christian Szegedy, Wei Liu, Yangqing Jia, Pierre Sermanet, Scott Reed,
Dragomir Anguelov, Dumitru Erhan, Vincent Vanhoucke, Andrew Rabinovich.
http://arxiv.org/pdf/1409.4842v1.pdf

import torch
import torch.nn as nn
from torchsummary import summary


# 卷积+激活,可重复利用
class BasicConv(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride=(1, 1), padding=(0, 0)):
        super(BasicConv, self).__init__()
        self.conv = nn.Conv2d(
            in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, stride=stride, padding=padding
        )
        self.relu = nn.ReLU(inplace=True)

    def forward(self, x):
        x = self.conv(x)
        x = self.relu(x)
        return x

# 推理模块
class InceptionBlock(nn.Module):
    def __init__(self, in_channels, ch1x1, ch3x3reduce, ch3x3, ch5x5reduce, ch5x5, chpool):
        super(InceptionBlock, self).__init__()
        self.branch_1 = BasicConv(in_channels=in_channels, out_channels=ch1x1, kernel_size=1)
        self.branch_2 = nn.Sequential(
            BasicConv(in_channels=in_channels, out_channels=ch3x3reduce, kernel_size=1),
            BasicConv(in_channels=ch3x3reduce, out_channels=ch3x3, kernel_size=3, padding=1)
        )
        self.branch_3 = nn.Sequential(
            BasicConv(in_channels=in_channels, out_channels=ch5x5reduce, kernel_size=1),
            BasicConv(in_channels=ch5x5reduce, out_channels=ch5x5, kernel_size=5, padding=2)
        )
        self.branch_4 = nn.Sequential(
            nn.MaxPool2d(kernel_size=3, padding=1, stride=1, ceil_mode=True),
            BasicConv(in_channels=in_channels, out_channels=chpool, kernel_size=1)
        )

    def forward(self, x):
        x_1 = self.branch_1(x)
        x_2 = self.branch_2(x)
        x_3 = self.branch_3(x)
        x_4 = self.branch_4(x)
        # 拼接
        x = torch.cat([x_1, x_2, x_3, x_4], dim=1)

        return x


class InceptionBlock_fc(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(InceptionBlock_fc, self).__init__()
        self.avg_pool = nn.AvgPool2d(kernel_size=5, stride=3)
        self.conv = BasicConv(in_channels=in_channels, out_channels=128, kernel_size=1, stride=1)
        # 全连接1 + 激活 + dropout
        self.auxiliary_linear1 = nn.Linear(in_features=128 * 4 * 4, out_features=1024)
        self.auxiliary_relu = nn.ReLU6(inplace=True)
        self.auxiliary_dropout = nn.Dropout(p=0.7)
        # 全连接2
        self.auxiliary_linear2 = nn.Linear(in_features=1024, out_features=out_channels)

    def forward(self, x):
        x = self.avg_pool(x)
        x = self.conv(x)
        x = x.view(x.size(0), -1)
        x = self.auxiliary_linear1(x)
        x = self.auxiliary_relu(x)
        x = self.auxiliary_dropout(x)
        x = self.auxiliary_linear2(x)
        return x


class Inception_V1(nn.Module):
    def __init__(self, num_classes):
        super(Inception_V1, self).__init__()
        self.BasicConv_1 = BasicConv(in_channels=3, out_channels=64, kernel_size=7, stride=2, padding=3)
        # ceil_mode用来决定不够滑动的区域是丢弃还是填充空值后运算
        self.max_pool_1 = nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True)
        self.conv_1x1 = BasicConv(in_channels=64, out_channels=64, kernel_size=1)
        self.conv_3x3 = BasicConv(in_channels=64, out_channels=192, kernel_size=3, padding=1)
        self.max_pool_2 = nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True)

        # in_channels, ch1x1, ch3x3reduce, ch3x3, ch5x5reduce, ch5x5, chpool
        self.InceptionBlock_3a = InceptionBlock(192, 64, 96, 128, 16, 32, 32)
        self.InceptionBlock_3b = InceptionBlock(256, 128, 128, 192, 32, 96, 64)

        self.max_pool_3 = nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True)

        self.InceptionBlock_4a = InceptionBlock(480, 192, 96, 208, 16, 48, 64)
        self.InceptionBlock_fc_1 = InceptionBlock_fc(in_channels=512, out_channels=num_classes)
        self.InceptionBlock_4b = InceptionBlock(512, 160, 112, 224, 24, 64, 64)
        self.InceptionBlock_4c = InceptionBlock(512, 128, 128, 256, 24, 64, 64)
        self.InceptionBlock_4d = InceptionBlock(512, 112, 144, 288, 32, 64, 64)
        self.InceptionBlock_4e = InceptionBlock(528, 256, 160, 320, 32, 128, 128)
        self.InceptionBlock_fc_2 = InceptionBlock_fc(in_channels=528, out_channels=num_classes)

        self.max_pool_4 = nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True)

        self.InceptionBlock_5a = InceptionBlock(832, 256, 160, 320, 32, 128, 128)
        self.InceptionBlock_5b = InceptionBlock(832, 384, 192, 384, 48, 128, 128)

        self.avg_pool = nn.AdaptiveAvgPool2d(1)  # 会把上一层的7x7压缩成1x1

        self.flatten = nn.Flatten()  # 降维打击
        self.fc = nn.Linear(in_features=1024, out_features=num_classes)

    def forward(self, x):
        x = self.BasicConv_1(x)
        x = self.max_pool_1(x)
        x = self.conv_1x1(x)
        x = self.conv_3x3(x)
        x = self.max_pool_2(x)

        x = self.InceptionBlock_3a(x)
        x = self.InceptionBlock_3b(x)

        x = self.max_pool_3(x)

        x = self.InceptionBlock_4a(x)
        x_1 = self.InceptionBlock_fc_1(x)

        x = self.InceptionBlock_4b(x)
        x = self.InceptionBlock_4c(x)
        x = self.InceptionBlock_4d(x)
        x_2 = self.InceptionBlock_fc_2(x)

        x = self.InceptionBlock_4e(x)
        x = self.max_pool_4(x)

        x = self.InceptionBlock_5a(x)
        x = self.InceptionBlock_5b(x)

        x = self.avg_pool(x)
        x = self.flatten(x)
        x = torch.dropout(x, 0.4, train=True)
        x = self.fc(x)

        x_3 = torch.softmax(x, dim=1)

        # x_1 = torch.softmax(x_1, dim=1)  # 在dim=1处做softmax操作
        return x_1, x_2, x_3

if __name__ == '__main__':
    # 只需要shape一致即可
    input = torch.ones([10, 3, 224, 224])

    # 调用模型
    model = Inception_V1(num_classes=20)

    # forward
    result = model(input)

    # 利用torch.summary观察每一层的状况
    summary(model.to("cuda"), (3, 224, 224))

    # torch.save(model.state_dict(), './model.pt')