LeNet 5 pytorch实现

import torch.nn as nn
import torch
from torchinfo import summary
#img is 32*32*1
data=torch.ones(size=(10,1,32,32))
class LeNet_5(nn.Module):
    def __init__(self):    #定义神经网络所需的全部元素
        super().__init__()
        self.conv1=nn.Conv2d(1,6,5)  #32+2pad-kernel  /1 +1
        self.pool1=nn.AvgPool2d(stride=2,kernel_size=2)
        self.conv2=nn.Conv2d(6,16,5)
        self.pool2=nn.AvgPool2d(stride=2,kernel_size=2)
        self.fc1=nn.Linear(5*5*16,120)
        self.fc2=nn.Linear(120,84)
        self.fc3=nn.Linear(84,10)

        #函数
        self.tanh=nn.Tanh()
        self.softmax=nn.Softmax(dim=1)

    def forward(self,x): #执行神经网络类，实际需要执行的运算
        x =self.tanh(self.conv1(x))
        x=self.pool1(x)
        x=self.tanh(self.conv2(x))
        x=self.pool2(x)
        #图像拉平
        x=x.view(-1,16*5*5)   #类似于reshape
        #-1:无论原来图像数据（sample）中是多少，在变形的时候和原来一样
        x=self.tanh(self.fc1(x))
        x = self.tanh(self.fc2(x))
        output=self.softmax(self.fc3(x))
        return output


net = LeNet_5()
print(net(data).shape)
#==net.forward(data)

#summary(net,input_size=(10,1,32,32))
summary(net)

结果：

=================================================================
Layer (type:depth-idx)                   Param #
=================================================================
LeNet_5                                  --
├─Conv2d: 1-1                            156
├─AvgPool2d: 1-2                         --
├─Conv2d: 1-3                            2,416
├─AvgPool2d: 1-4                         --
├─Linear: 1-5                            48,120
├─Linear: 1-6                            10,164
├─Linear: 1-7                            850
├─Tanh: 1-8                              --
├─Softmax: 1-9                           --
=================================================================
Total params: 61,706
Trainable params: 61,706
Non-trainable params: 0
=================================================================