pytorch入门之MNIST数据集分类

本着入门pytorch的热情，实现pytorch的"helloword"

概述：搭建LeNet-5卷积神经网络，该网络具体参数配置见代码。代码详解，代码有详细注释：ps:我这是用GPU跑的代码，如果要用cpu跑，则需要把train中的.cuda()一行删除，并适当降低卷积神经元的个数。

mycnn部分：

import torch.nn as nn class MyCNN(nn.Module): # 继承module def __init__(self): # 定义构造函数 super(MyCNN, self).__init__() # 计算公式pading=(k_size-1)/2 self.con1 = nn.Sequential( nn.Conv2d(1, 6, 5, 1, 2), # 1*28*28--->=6*28*28 nn.MaxPool2d(2), # 6*28*28-->=6*14*14 nn.ReLU() ) self.con2 = nn.Sequential( nn.Conv2d(6, 60, 5, 1, 2), # 6*14*14-->60*14*14 nn.MaxPool2d(2), # 60*14*14-->60*7*7 nn.ReLU() ) self.con3 = nn.Sequential( nn.Conv2d(60, 96, 5, 1, 2), # 60*7*7-->96*7*7 nn.ReLU() ) self.fc = nn.Sequential( nn.Linear(96 * 7 * 7, 84), nn.ReLU(), nn.Linear(84, 10) ) def forward(self, inputs): out = self.con1(inputs) out = self.con2(out) out = self.con3(out) out = out.view(out.size(0), -1) out = self.fc(out) return out

load_data部分

import torch from torchvision import datasets, transforms from torch.utils.data import DataLoader # 下载训练集 train_dataset = datasets.MNIST(root='./num/', train=True, transform=transforms.ToTensor(), download=True) # 下载测试集 test_dataset = datasets.MNIST(root='./num/', train=False, transform=transforms.ToTensor(), download=True) # dataset 参数用于指定我们载入的数据集名称 # batch_size参数设置了每个包中的图片数据个数 # 在装载的过程会将数据随机打乱顺序并进打包 # 建立一个数据迭代器 # 装载训练集 train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=64, shuffle=True) # 装载测试集 test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=64, shuffle=True)

train部分：

import torch import torch.nn as nn import torch.optim as optim from torch.autograd import Variable from mycnn import MyCNN from load_data import test_dataset, train_loader # 如果有GPU就用GPU否则使用cpu device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # batch_size = 64 # 定义反向传播学习率 LR = 0.001 # 实例化卷积类 net = MyCNN().to(device) # 损失函数使用交叉熵 criterion = nn.CrossEntropyLoss() # 优化函数使用 Adam 自适应优化算法 optimizer = optim.Adam( net.parameters(), lr=LR, ) # 批次数 epoch = 1 if __name__ == '__main__': for epoch in range(epoch): # 总体损失为0 sum_loss = 0.0 # enumerate() 函数用于将一个可遍历的数据对象(如列表、元组或字符串)组合为一个索引序列，同时列出数据和数据下标 for i, data in enumerate(train_loader): inputs, labels = data # 使用GPU时候使用 inputs, labels = Variable(inputs).cuda(), Variable(labels).cuda() # 将梯度归零 optimizer.zero_grad() # 将数据传入网络进行前向运算 outputs = net(inputs) # 得到损失函数 loss = criterion(outputs, labels) # 反向传播 loss.backward() # 通过梯度做一步参数更新 optimizer.step() # print(loss) sum_loss += loss.item() if i % 100 == 99: print('[%d,%d] loss:%.03f' % (epoch + 1, i + 1, sum_loss / 100)) sum_loss = 0.0

test部分：

from train import net from load_data import test_dataset, test_loader import torch from torch.autograd import Variable net.eval() # 将模型变换为测试模式 correct = 0 total = 0 for data_test in test_loader: images, labels = data_test images, labels = Variable(images).cuda(), Variable(labels).cuda() output_test = net(images) _, predicted = torch.max(output_test, 1) # predicted为返回的最大的概率值的索引 total += labels.size(0) correct += (predicted == labels).sum() print("correct1: ", correct) print("Test acc: {0}".format(correct.item() / len(test_dataset)))

工程文件：github 参考文献：特别鸣谢