网络结构
LeNet现在主要指LeNet5,主要特征是将卷积层和下采样层相结合作为网络的基本结构。
输入为一个矩阵或者图像,大小为32X32。不计输入层,这个模型共有7层,3个卷积层,2个下采样层,1个全连接层和1个输出层。
C1:第一个卷积层。包含6个卷积特征图,每个特征图大小为28X28,由一个5X5的卷积核对输入图像进 行内卷积运算得到。
S2:第一个下采样层(池化层)。包含6个14X14的下采样特征图。每个下采样特征图都是由C1层的特征图经过2X2,步长为2的窗口进行平均池化,再利用sigmoid进行非线性变换处理。
C3:第二个卷积层。包含16个10X10卷积特征图。由一个5X5的卷积核对输入图像进行内卷积运算得到。
S4:第一个下采样层(池化层)。包含16个5X5的下采样特征图。每个下采样特征图都是由C1层的特征图经过2X2,步长为2的窗口进行平均池化,再利用sigmoid进行非线性变换处理。
C5:第三个卷积层。包含120个1X1卷积特征图。由一个5X5的卷积核对输入图像进行内卷积运算得到。
F6:全连接层。包含84个节点。
最后一层为输出层。全连接层。
网络代码
from torch import nn
class LeNet(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Sequential(
nn.Conv2d(3, 6, kernel_size=5, stride=1),
nn.ReLU())
self.max_pool1 = nn.MaxPool2d(2)
self.conv2 = nn.Sequential(nn.Conv2d(6, 16, kernel_size=5), nn.ReLU())
self.max_pool2 = nn.MaxPool2d(2)
self.fc1 = nn.Sequential(nn.Linear(16 * 5 * 5, 120), nn.ReLU())
self.fc2 = nn.Sequential(nn.Linear(120, 84), nn.ReLU())
self.fc3 = nn.Linear(84, 10)
def forward(self, x):
x = self.conv1(x)
x = self.max_pool1(x)
x = self.conv2(x)
x = self.max_pool2(x)
x = x.view(x.shape[0], -1)
x = self.fc1(x)
x = self.fc2(x)
x = self.fc3(x)
return x
接下来进行训练,代码如下
训练代码
import torch
import torchvision
import torch.nn as nn
import torch.optim as optim
import torchvision.transforms as transforms
from model import LeNet
from torch.autograd import Variable
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# 准备数据集并进行预处理
transform = transforms.Compose([transforms.ToTensor