import torch
from torch import nn

# 建神经网络
class Tudui(nn.Module):
    def __init__(self):
        super(Tudui, self).__init__()
        self.model = nn.Sequential(
            nn.Conv2d(3, 32, 5, 1, 2),  # 输入通道数为3，输出通道数为32，卷积核大小为5x5，步长为1，填充为2
            nn.MaxPool2d(2),  # 最大池化层，池化窗口大小为2x2
            nn.Conv2d(32, 32, 5, 1, 2),  # 输入通道数为32，输出通道数为32，卷积核大小为5x5，步长为1，填充为2
            nn.MaxPool2d(2),  # 最大池化层，池化窗口大小为2x2
            nn.Conv2d(32, 64, 5, 1, 2),  # 输入通道数为32，输出通道数为64，卷积核大小为5x5，步长为1，填充为2
            nn.MaxPool2d(2),  # 最大池化层，池化窗口大小为2x2
            nn.Flatten(),  # 将多维张量展平为一维
            nn.Linear(64 * 4 * 4, 64),  # 全连接层，输入节点数为64 * 4 * 4，输出节点数为64
            nn.Linear(64, 10)  # 连接层，输入节点数为64，输出节点数为10
        )

    def forward(self, x):
        x = self.model(x)
        return x

if __name__ == '__main__':
    tudui = Tudui()
    input = torch.ones((64, 3, 32, 32))
    output = tudui(input)
    print(output.shape)