首页
学习
活动
专区
工具
TVP
发布
精选内容/技术社群/优惠产品,尽在小程序
立即前往

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元无门槛券

私享最新 技术干货

扫码加入开发者社群
领券