复现InceptionV2+V3

论文地址：

https://www.cv-foundation.org/openaccess/content_cvpr_2016/papers/Szegedy_Rethinking_the_Inception_CVPR_2016_paper.pdf

其中InceptionA/B/C分别代表论文中的三个重复使用的结构

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 InceptionA(nn.Module):
    # pool_channels用来控制使池化输出为288的通道--复现论文要求
    def __init__(self, in_channels, pool_channels):
        super(InceptionA, self).__init__()
        self.branch_1 = nn.Sequential(
            nn.Conv2d(in_channels=in_channels, out_channels=64, kernel_size=1),
            nn.Conv2d(in_channels=64, out_channels=96, kernel_size=3, padding=1),
            nn.Conv2d(in_channels=96, out_channels=96, kernel_size=3, padding=1)
        )

        self.branch_2 = nn.Sequential(
            nn.Conv2d(in_channels=in_channels, out_channels=48, kernel_size=1),
            nn.Conv2d(in_channels=48, out_channels=64, kernel_size=3, padding=1)
        )

        self.branch_3 = nn.Sequential(
            nn.MaxPool2d(kernel_size=3, stride=1, padding=1),
            nn.Conv2d(in_channels=in_channels, out_channels=pool_channels, kernel_size=1)
        )

        self.branch_4 = nn.Conv2d(in_channels=in_channels, out_channels=64, 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)

        # 4个分支输出的shape相同，可用cat在dim=1处拼接，96+64+64+64=288符合输出
        # print("x_1.shape=", x_1.shape)
        # print("x_2.shape=", x_2.shape)
        # print("x_3.shape=", x_3.shape)
        # print("x_4.shape=", x_4.shape)

        # 拼接4个分支
        x = torch.cat([x_1, x_2, x_3, x_4], dim=1)
        return x

class InceptionB(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(InceptionB, self).__init__()
        out_channels = out_channels // 4  # 为了控制拼接4个branch
        self.branch_1 = nn.Sequential(
            BasicConv(in_channels=in_channels, out_channels=out_channels, kernel_size=(1)),
            BasicConv(in_channels=out_channels, out_channels=out_channels, kernel_size=(1, 7), padding=(0, 3)),
            BasicConv(in_channels=out_channels, out_channels=out_channels, kernel_size=(7, 1), padding=(3, 0)),
            BasicConv(in_channels=out_channels, out_channels=out_channels, kernel_size=(1, 7), padding=(0, 3)),
            BasicConv(in_channels=out_channels, out_channels=out_channels, kernel_size=(7, 1), padding=(3, 0))
        )
        self.branch_2 = nn.Sequential(
            BasicConv(in_channels=in_channels, out_channels=out_channels, kernel_size=1),
            BasicConv(in_channels=out_channels, out_channels=out_channels, kernel_size=(1, 7), padding=(0, 3)),
            BasicConv(in_channels=out_channels, out_channels=out_channels, kernel_size=(7, 1), padding=(3, 0))
        )
        self.branch_3 = nn.Sequential(
            nn.MaxPool2d(kernel_size=3, stride=1, padding=1),
            BasicConv(in_channels=in_channels, out_channels=out_channels, kernel_size=1, stride=1)
        )
        self.branch_4 = BasicConv(in_channels=in_channels, out_channels=out_channels, 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)

        # 检验
        # print("x_1.shape=", x_1.shape)
        # print("x_2.shape=", x_2.shape)
        # print("x_3.shape=", x_3.shape)
        # print("x_4.shape=", x_4.shape)
        x = torch.cat([x_1, x_2, x_3, x_4], dim=1)
        return x

class InceptionC(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(InceptionC, self).__init__()
        out_channels = out_channels//6  # 是为了把6个branch拼接起来
        self.branch_1_1 = nn.Sequential(
            BasicConv(in_channels=in_channels, out_channels=out_channels, kernel_size=1),
            BasicConv(in_channels=out_channels, out_channels=out_channels, kernel_size=3, padding=1)
        )
        self.branch_1_2_1 = BasicConv(in_channels=out_channels, out_channels=out_channels, kernel_size=(1, 3), padding=(0, 1))
        self.branch_1_2_2 = BasicConv(in_channels=out_channels, out_channels=out_channels, kernel_size=(3, 1), padding=(1, 0))

        self.branch_2_1 = BasicConv(in_channels=in_channels, out_channels=out_channels, kernel_size=1)
        self.branch_2_2_1 = BasicConv(in_channels=out_channels, out_channels=out_channels, kernel_size=(1, 3), padding=(0, 1))
        self.branch_2_2_2 = BasicConv(in_channels=out_channels, out_channels=out_channels, kernel_size=(3, 1), padding=(1, 0))

        self.branch_3 = nn.Sequential(
            nn.MaxPool2d(kernel_size=3, stride=1, padding=1),
            BasicConv(in_channels=in_channels, out_channels=out_channels+1, kernel_size=1)  # 为了凑论文通道--out_channels//6凑数取整出的问题
        )

        self.branch_4 = BasicConv(in_channels=in_channels, out_channels=out_channels+1, kernel_size=1)  # 为了凑论文通道--out_channels//6凑数取整出的问题

    def forward(self, x):
        x_1_1 = self.branch_1_1(x)
        x_1_2_1 = self.branch_1_2_1(x_1_1)
        x_1_2_2 = self.branch_1_2_2(x_1_1)

        x_2_1 = self.branch_2_1(x)
        x_2_2_1 = self.branch_2_2_1(x_2_1)
        x_2_2_2 = self.branch_2_2_2(x_2_1)

        x_3 = self.branch_3(x)

        x_4 = self.branch_4(x)

        # 检验
        # print("x_1_2_1.shape=", x_1_2_1.shape)
        # print("x_1_2_2.shape=", x_1_2_2.shape)
        # print("x_2_2_1.shape=", x_2_2_1.shape)
        # print("x_2_2_2.shape=", x_2_2_2.shape)
        # print("x_3.shape=", x_3.shape)
        # print("x_4.shape=", x_4.shape)
        x = torch.cat([x_1_2_1, x_1_2_2, x_2_2_1, x_2_2_2, x_3, x_4], dim=1)
        return x

class Model_Expand(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(Model_Expand, self).__init__()
        self.branch_1 = nn.Sequential(
            BasicConv(in_channels=in_channels, out_channels=out_channels, kernel_size=1),
            BasicConv(in_channels=out_channels, out_channels=out_channels, kernel_size=3, stride=1),
            BasicConv(in_channels=out_channels, out_channels=out_channels, kernel_size=3, stride=2, padding=1)
        )
        self.branch_2 = nn.Sequential(
            BasicConv(in_channels=in_channels, out_channels=out_channels, kernel_size=1),
            BasicConv(in_channels=out_channels, out_channels=out_channels, kernel_size=3, stride=2)
        )
        self.branch_3 = nn.MaxPool2d(kernel_size=3, stride=2)

    def forward(self, x):
        x_1 = self.branch_1(x)
        x_2 = self.branch_2(x)
        x_3 = self.branch_3(x)

        # # 检验
        # print("x_1.shape=", x_1.shape)
        # print("x_2.shape=", x_2.shape)
        # print("x_3.shape=", x_3.shape)

        x = torch.cat([x_1, x_2, x_3], dim=1)
        return x

class Inception_V2(nn.Module):
    def __init__(self, num_classes):
        super(Inception_V2, self).__init__()

        self.conv_1 = BasicConv(in_channels=3, out_channels=32, kernel_size=3, stride=2)
        self.conv_2 = BasicConv(in_channels=32, out_channels=32, kernel_size=3, stride=1)
        self.conv_padded = BasicConv(in_channels=32, out_channels=64, kernel_size=3, stride=1, padding=1)
        self.pool_1 = nn.MaxPool2d(kernel_size=3, stride=2)
        self.conv_3 = BasicConv(in_channels=64, out_channels=80, kernel_size=3, stride=1)
        self.conv_4 = BasicConv(in_channels=80, out_channels=192, kernel_size=3, stride=2)
        self.conv_5 = BasicConv(in_channels=192, out_channels=288, kernel_size=3, stride=1, padding=1)

        self.Inception_a_1 = InceptionA(in_channels=288, pool_channels=64)
        self.Inception_a_2 = InceptionA(in_channels=288, pool_channels=64)
        self.Inception_a_3 = InceptionA(in_channels=288, pool_channels=64)

        self.Module_Expand_1 = Model_Expand(288, 240)

        self.Inception_b_1 = InceptionB(768, 768)
        self.Inception_b_2 = InceptionB(768, 768)
        self.Inception_b_3 = InceptionB(768, 768)
        self.Inception_b_4 = InceptionB(768, 768)

        # 注意这里控制输出的是256通道,但实际有三个分支,总通道数=256+256+768=1280匹配论文
        self.Module_Expand_2 = Model_Expand(768, 256)

        self.Inception_c_1 = InceptionC(1280, 1280)
        self.Inception_c_2 = InceptionC(1280, 2048)

        self.pool_2 = nn.AdaptiveAvgPool2d(1)

        self.flatten = nn.Flatten()

        self.fc_1 = nn.Linear(in_features=2048, out_features=num_classes)  # 输出为类别数

    def forward(self, x):
        x = self.conv_1(x)
        x = self.conv_2(x)
        x = self.conv_padded(x)
        x = self.pool_1(x)
        x = self.conv_3(x)
        x = self.conv_4(x)
        x = self.conv_5(x)

        x = self.Inception_a_1(x)
        x = self.Inception_a_2(x)
        x = self.Inception_a_3(x)

        x = self.Module_Expand_1(x)
        x = self.Inception_b_1(x)
        x = self.Inception_b_2(x)
        x = self.Inception_b_3(x)
        x = self.Inception_b_4(x)

        x = self.Module_Expand_2(x)

        x = self.Inception_c_1(x)
        x = self.Inception_c_2(x)

        x = self.pool_2(x)
        x = self.flatten(x)
        x = self.fc_1(x)
        x = torch.softmax(x, dim=1)
        return x

if __name__ == '__main__':
    # 根据第一层的输入要求来设定,第一个参数表示共10个branch
    input = torch.ones([10, 3, 299, 299])
    model = Inception_V2(num_classes=5)  # 共5类
    output = model(input)
    summary(model.to("cuda"), (3, 299, 299))
    print(output.shape)  # 用输出尺寸来判断是否正确复现论文模型

InceptionV3和V2的区别就是在BasicConv中添加了一层bn层

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.bn = nn.BatchNorm2d(num_features=out_channels)  # 与V2的区别
        self.relu = nn.ReLU(inplace=True)

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