Как добавить слой свертки между двумя другими сверточными слоями?

Принцип такой - число входных каналов в новом промежуточном слое должно равняться числу выходных каналов предыдущего слоя, а число выходных каналов в новом слое равняться числу входных следующего слоя:

class LeNet5(torch.nn.Module):
    def __init__(self):
        super(LeNet5, self).__init__()

        self.conv1 = torch.nn.Conv2d(
            in_channels=1, out_channels=6, kernel_size=5, padding=2)
        #######
        self.conv1_1 = torch.nn.Conv2d(
            in_channels=6, out_channels=6, kernel_size=5, padding=2)
        #######
        self.act1  = torch.nn.Tanh()
        self.pool1 = torch.nn.AvgPool2d(kernel_size=2, stride=2)

        self.pool1 = torch.nn.AvgPool2d(kernel_size=2, stride=2)


        self.conv2 = torch.nn.Conv2d(
            in_channels=6, out_channels=16, kernel_size=5, padding=0)
        self.act2  = torch.nn.Tanh()
        self.pool2 = torch.nn.AvgPool2d(kernel_size=2, stride=2)

        self.fc1   = torch.nn.Linear(5 * 5 * 16, 120)
        self.act3  = torch.nn.Tanh()

        self.fc2   = torch.nn.Linear(120, 84)
        self.act4  = torch.nn.Tanh()

        self.fc3   = torch.nn.Linear(84, 10)

    def forward(self, x):

        x = self.conv1(x)
        x = self.act1(x)
        #######
        x = self.conv1_1(x)
        x = self.act1(x)
        #######
        x = self.pool1(x)

        x = self.conv2(x)
        x = self.act2(x)
        x = self.pool2(x)

        x = x.view(x.size(0), x.size(1) * x.size(2) * x.size(3))

        x = self.fc1(x)
        x = self.act3(x)
        x = self.fc2(x)
        x = self.act4(x)
        x = self.fc3(x)

        return x

как вариант, делайте ошибки. нестыковки числа выходных и входных слоев при запуске скрипта дают исключения с указанием чисел, которые и надо совместить.

например,

RuntimeError: running_mean should contain 768 elements not 32

для

   nn.Conv2d(48, **768**, 2, 1, 1), nn.BatchNorm2d(**32**), nn.ReLU(True),  

или

RuntimeError: Given groups=1, weight of size [6, 512, 2, 2], expected input[300, 768, 33, 9] to have 512 channels, but got 768 channels instead

    nn.Conv2d(48, 768, 2, 1, 1), nn.BatchNorm2d(768), nn.ReLU(True), 

    nn.Conv2d(512, 6, 2, 1), nn.ReLU(True), nn.MaxPool2d(2, 2),

на SO еще часто ссылаются на формулу, учитывающую каналы, пэддинг и страйдинг,

также используют визуализацию модели

ZZNet(
  (cnn): Sequential(
    (0): Conv2d(3, 3, kernel_size=(1, 1), stride=(1, 1), padding=(1, 1))
    (1): ReLU(inplace=True)
    (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (3): Conv2d(3, 48, kernel_size=(2, 2), stride=(1, 1))
    (4): ReLU(inplace=True)
    (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (6): Conv2d(48, 768, kernel_size=(2, 2), stride=(1, 1), padding=(1, 1))
    (7): BatchNorm2d(768, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (8): ReLU(inplace=True)
    (9): Conv2d(512, 6, kernel_size=(2, 2), stride=(1, 1))
    (10): ReLU(inplace=True)
    (11): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (12): Conv2d(6, 6, kernel_size=(2, 2), stride=(1, 1), padding=(1, 1))
    (13): BatchNorm2d(6, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (14): ReLU(inplace=True)
    (15): Conv2d(6, 6, kernel_size=(2, 2), stride=(1, 1))
    (16): ReLU(inplace=True)
    (17): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (18): Conv2d(6, 6, kernel_size=(2, 2), stride=(1, 1), padding=(1, 1))
    (19): BatchNorm2d(6, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (20): ReLU(inplace=True)
    (21): Conv2d(6, 6, kernel_size=(2, 2), stride=(1, 1))
    (22): ReLU(inplace=True)
    (23): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (24): Conv2d(6, 6, kernel_size=(2, 2), stride=(1, 1), padding=(1, 1))
    (25): BatchNorm2d(6, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (26): ReLU(inplace=True)
    (27): LeakyReLU(negative_slope=0.1)
    (28): Flatten()
    (29): Linear(in_features=60, out_features=60, bias=True)
    (30): LeakyReLU(negative_slope=0.1)
    (31): Linear(in_features=60, out_features=12, bias=True)
  )
)

update, визуализация слоев

net = ZZNet().to(device)

#1
print(net)

#2
# Print model's state_dict
print("Model's state_dict:")
for param_tensor in net.state_dict():
    print(param_tensor, "\t", net.state_dict()[param_tensor].size())


#3
torchstat.stat(net, (3, 32, 128))

#4 - отброшенные варианты
# for idx, m in enumerate(net.modules()):
#         print(idx, '->', m, type(m))
#
# # for name, param in net.parameters():
# #     if name in ['bias']:
# #         print('param: ', param.size())
#
# print('NN: ',net.parameters())
# for param in net.parameters():
#     print(type(param.data), param.size())