Pytorch小试牛刀（二）——LeNet
训练数据集：CIFAR10
没有对该模型做任何改进，完全按照论文要求进行编码。
LeNet-5结构（很简单，在这里不做讲解）：
缺点：LeNet-5采用的激活函数为Sigmoid激活函数，会产生梯度消失的现象。
代码如下：
import torch import torch.nn as nnimport torchvisionimport torch.nn.functional as Fimport torchvision.transforms as transformsfrom torch.autograd import Variable
'''定义参数'''batch_size=64lr=0.001num_classes=10
'''获取数据集'''train_dataset=torchvision.datasets.CIFAR10(root='../DataSet/',train=True,download=False,transform=transforms.ToTensor())
test_dataset=torchvision.datasets.CIFAR10(root='../DataSet/',train=False,download=False,transform=transforms.ToTensor())
'''装载数据'''train_loader=torch.utils.data.DataLoader(dataset=train_dataset,batch_size=batch_size,shuffle=True)test_loader=torch.utils.data.DataLoader(dataset=test_dataset,batch_size=batch_size,shuffle=False)
'''设计LeNet_5'''class LeNet(nn.Module):def __init__(self,num_classes=10): super(LeNet, self).__init__()'''第一层卷积，卷积核大小为5*5，步距为1，输入通道为3，输出通道为6''' self.conv1=nn.Conv2d(3,6,kernel_size=5,stride=1)
'''第一层池化层，卷积核为2*2，步距为2，相当于特征图缩小了一般''' self.pool1=nn.MaxPool2d(kernel_size=2,stride=2)
'''第二层卷积，卷积核大小为5*5，步距为1，输入通道为6，输出通道为16''' self.conv2=nn.Conv2d(6,16,kernel_size=5,stride=1)
'''第二层池化层，卷积核为2*2，步距为2，相当于特征图缩小了一般''' self.pool2=nn.MaxPool2d(kernel_size=2,stride=2)
'''第一层全连接层，维度由16*5*5=>120''' self.linear1=nn.Linear(16*5*5,120)
'''第二层全连接层，维度由120=>84''' self.linear2=nn.Linear(120,84)
'''第三层全连接层，维度由64=>10''' self.linear3=nn.Linear(84,num_classes)
def forward(self, x):'''将数据送入第一个卷积层''' out=torch.sigmoid(self.conv1(x))
'''将数据送入第一个池化层''' out=self.pool1(out)
'''将数据送入第二个卷积层''' out=torch.sigmoid(self.conv2(out))
'''将数据送入第二个池化层''' out=self.pool2(out)
'''将池化层后的数据进行Flatten，使数据变成能够被FC层接受的Vector''' out=out.reshape(-1,16*5*5)
'''将数据送入第一个全连接层''' out=torch.sigmoid(self.linear1(out))
'''将数据送入第二个全连接层''' out=torch.sigmoid(self.linear2(out))
'''将数据送入第三个全连接层得到输出''' out=self.linear3(out)
return out
model = LeNet(num_classes)
'''设置损失函数'''criterion=nn.CrossEntropyLoss()
'''设置优化器'''optimizer=torch.optim.Adam(model.parameters(),lr=lr)
'''开始训练'''total_step = len(train_loader)for epoch in range(10):for i, (images, labels) in enumerate(train_loader):
images=Variable(images) labels=Variable(labels)# Forward pass outputs = model(images) loss = criterion(outputs, labels)
# Backward and optimize optimizer.zero_grad() loss.backward() optimizer.step()
if (i+1) % 100 == 0:print ('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}' .format(epoch+1, 10, i+1, total_step, loss.item()))
结果如下：
- 发表于:
- 原文链接：https://kuaibao.qq.com/s/20200921A0EL0L00?refer=cp_1026
- 腾讯「腾讯云开发者社区」是腾讯内容开放平台帐号（企鹅号）传播渠道之一，根据《腾讯内容开放平台服务协议》转载发布内容。
- 如有侵权，请联系 cloudcommunity@tencent.com 删除。
相关快讯
扫码
添加站长 进交流群
领取专属 10元无门槛券
私享最新 技术干货
腾讯云开发者
