PyTorch搭建LeNet神经网络

news2024/11/18 21:37:42

函数的参数

1、PyTorch Tensor的通道排序

[batch, channel, height, width]

batch: 要处理的一批图像的个数

channel: 通道数（一般是R G B 三个通道）

height: 图像的高度

width: 图像的宽度

2.Conv 2d 卷积层的参数

[in_channels, out_channels, kernel_size, stride = 1,padding = 0]

in_channels: 输入特征矩阵的深度（一般为RGB深度为3）

out_channels:使用卷积核的个数（使用n个卷积核，那么就会生成一个深度为n维的特征矩阵）

kernel_size:卷积和的大小

stride = 1: 默认步长为1

padding = 0: 默认填充为0

卷积层输出尺寸计算公式看深度学习-图像处理-卷积神经网络_cnn卷积神经网络像素处理-CSDN博客

3.MaxPool2d 池化层的参数

[kernel_size, stride = None, padding = 0]

kernel_size:池化核大小，手动设定

stride = None:步长不设置默认值，但如果没有指定步长，将会与池化核大小相同

padding = 0:默认填充为0

池化层不改变深度，只改变图像的高度宽度

4.全连接层

全连接层需要输入一维向量，所有我们需要把前面得到的向量展平，通过使用view函数把向量数据展成一维向量。

通过设置节点个数来缩小数据，第一层设置为120，第二层设置为84，第三次最后的设置大小根据数据自身的情况类别来设置。

代码实现

import torch.nn as nn   # 导入包
import torch.nn.functional as F


class LeNet(nn.Module):
    def __init__(self):  # 初始化函数
        super(LeNet, self).__init__()
        self.conv1 = nn.Conv2d(3, 16, 5)  # 卷积层
        self.pool1 = nn.MaxPool2d(2, 2)  # 下采样层
        self.conv2 = nn.Conv2d(16, 32, 5)
        self.pool2 = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(32*5*5, 120)  # 全链接层
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):  # 正向传播函数
        x = F.relu(self.conv1(x))  # 采用relu激活函数
        x = self.pool1(x)
        x = F.relu(self.conv2(x))
        x = self.pool2(x)
        x = x.view(-1, 32*5*5)  # view函数实现数据展成一维向量
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

# 测试实例
import torch
input1 = torch.rand([32, 3, 32, 32])  # 输入tensor的通道参数
mode1 = LeNet()
print(mode1)
output = mode1(input1)

运行结果