class ConvReLU(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, stride=stride, padding=padding, bias=True),
            nn.ReLU(inplace=True),
        )    
    def forward(self, x):
        return self.conv(x)
        
        
class InceptionModule(nn.Module):
    def __init__(self, in_channels, c1x1_out, c3x3_in, c3x3_out, c5x5_in, c5x5_out, pool_proj):
        super().__init__()
        self.branch1 = ConvReLU(in_channels=in_channels, out_channels=c1x1_out, kernel_size=1, stride=1, padding=0)
        self.branch2 = nn.Sequential(
            ConvReLU(in_channels=in_channels, out_channels=c3x3_in, kernel_size=1, stride=1, padding=0),
            ConvReLU(in_channels=c3x3_in, out_channels=c3x3_out, kernel_size=3, stride=1, padding=1)
        )
        self.branch3 = nn.Sequential(
            ConvReLU(in_channels=in_channels, out_channels=c5x5_in, kernel_size=1, stride=1, padding=0),
            ConvReLU(in_channels=c5x5_in, out_channels=c5x5_out, kernel_size=5, stride=1, padding=2)
        )
        self.branch4 = nn.Sequential(
            nn.MaxPool2d(kernel_size=3, stride=1, padding=1),
            ConvReLU(in_channels=in_channels, out_channels=pool_proj, kernel_size=1, stride=1, padding=0)
        )
    def forward(self, x):
        x1 = self.branch1(x)
        x2 = self.branch2(x)
        x3 = self.branch3(x)
        x4 = self.branch4(x)
        x = torch.cat([x1, x2, x3, x4], dim=1)
        return x



class AuxClassifier(nn.Module):
    def __init__(self, in_channels, n_classes):
        super().__init__()
        self.avgpool = nn.AdaptiveAvgPool2d(4)
        self.conv = ConvReLU(in_channels=in_channels, out_channels=128, kernel_size=1, stride=1, padding=0)
        self.fc1 = nn.Sequential(
            nn.Linear(in_features=128*4*4, out_features=1024, bias=True),
            nn.ReLU(inplace=True)
        )
        self.dropout = nn.Dropout(p=0.7)
        self.fc2 = nn.Linear(in_features=1024, out_features=n_classes, bias=True)
        self.softmax = nn.Softmax(dim=-1)
        
    def forward(self, x):
        b, _, _ ,_ = x.shape
        x = self.avgpool(x)
        x = self.conv(x)
        x = self.fc1(x.view(b, -1))
        x = self.dropout(x)
        x = self.fc2(x)
        x = self.softmax(x)
        return x
    
    
class GooLeNet(nn.Module):
    def __init__(self, in_channels, n_classes) -> None:
        super().__init__()
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        self.avgpool = nn.AdaptiveAvgPool2d(output_size=1)
        
        self.conv1 = nn.Sequential(
            ConvReLU(in_channels=in_channels, out_channels=64, kernel_size=7, stride=2, padding=3),
            nn.LocalResponseNorm(size=5, k=2, alpha=1e-4, beta=0.75),
        )
        self.conv2 = nn.Sequential(
            ConvReLU(in_channels=64, out_channels=64, kernel_size=1, stride=1, padding=0),
            ConvReLU(in_channels=64, out_channels=192, kernel_size=3, stride=1, padding=1),
            nn.LocalResponseNorm(size=5, k=2, alpha=1e-4, beta=0.75),
        )
        
        self.inception3a = InceptionModule(in_channels=192, c1x1_out=64, c3x3_in=96, c3x3_out=128, c5x5_in=16, c5x5_out=32, pool_proj=32)
        self.inception3b = InceptionModule(in_channels=256, c1x1_out=128, c3x3_in=128, c3x3_out=192, c5x5_in=32, c5x5_out=96, pool_proj=64)

        self.inception4a = InceptionModule(in_channels=480, c1x1_out=192, c3x3_in=96, c3x3_out=208, c5x5_in=16, c5x5_out=48, pool_proj=64)
        self.inception4b = InceptionModule(in_channels=512, c1x1_out=160, c3x3_in=112, c3x3_out=224, c5x5_in=24, c5x5_out=64, pool_proj=64)
        self.inception4c = InceptionModule(in_channels=512, c1x1_out=128, c3x3_in=128, c3x3_out=256, c5x5_in=24, c5x5_out=64, pool_proj=64)
        self.inception4d = InceptionModule(in_channels=512, c1x1_out=112, c3x3_in=144, c3x3_out=288, c5x5_in=32, c5x5_out=64, pool_proj=64)
        self.inception4e = InceptionModule(in_channels=528, c1x1_out=256, c3x3_in=160, c3x3_out=320, c5x5_in=32, c5x5_out=128, pool_proj=128)

        self.inception5a = InceptionModule(in_channels=832, c1x1_out=256, c3x3_in=160, c3x3_out=320, c5x5_in=32, c5x5_out=128, pool_proj=128)
        self.inception5b = InceptionModule(in_channels=832, c1x1_out=384, c3x3_in=192, c3x3_out=384, c5x5_in=48, c5x5_out=128, pool_proj=128)

        self.dropout = nn.Dropout(p=0.4)
        self.fc = nn.Linear(in_features=1024, out_features=n_classes, bias=True)
        self.softmax = nn.Softmax(dim=-1)
        
        self.aux_classfier1 = AuxClassifier(in_channels=512, n_classes=n_classes)
        self.aux_classfier2 = AuxClassifier(in_channels=528, n_classes=n_classes)

    def forward(self, x):
        b, _, _, _ = x.shape
        x = self.conv1(x)
        print('# Conv1 output shape:', x.shape)
        x = self.maxpool(x)
        print('# Pool1 output shape:', x.shape)
        x = self.conv2(x)
        print('# Conv2 output shape:', x.shape)
        x = self.maxpool(x)
        print('# Pool2 output shape:', x.shape)
        x = self.inception3a(x)
        print('# Inception3a output shape:', x.shape)
        x = self.inception3b(x)
        print('# Inception3b output shape:', x.shape)
        x = self.maxpool(x)
        print('# Pool3 output shape:', x.shape)
        x = self.inception4a(x)
        print('# Inception4a output shape:', x.shape)
        aux1 = self.aux_classfier1(x)
        print('# aux_classifier1 output shape:', aux1.shape)
        x = self.inception4b(x)
        print('# Inception4b output shape:', x.shape)
        x = self.inception4c(x)
        print('# Inception4c output shape:', x.shape)
        x = self.inception4d(x)
        print('# Inception4d output shape:', x.shape)
        aux2 = self.aux_classfier2(x)
        print('# aux_classifier2 output shape:', aux2.shape)
        x = self.inception4e(x)
        print('# Inception4e output shape:', x.shape)
        x = self.maxpool(x)
        print('# Pool4 output shape:', x.shape)
        x = self.inception5a(x)
        print('# Inception5a output shape:', x.shape)
        x = self.inception5b(x)
        print('# Inception5b output shape:', x.shape)
        x = self.avgpool(x)
        print('# Avgpool output shape:', x.shape)
        x = self.dropout(x.view(b, -1))
        print('# dropout output shape:', x.shape)
        x = self.fc(x)
        print('# FC output shape:', x.shape)
        x = self.softmax(x)
        print('# Softmax output shape:', x.shape)
        return x, aux1, aux2

    
inputs = torch.randn(4, 3, 224, 224)
cnn = GooLeNet(in_channels = 3, n_classes = 1000)
outputs = cnn(inputs)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146