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
=================================================================