MNIST 手写数字识别
- 导入相关库
import torch
# pytorch中最重要的模块，封装了神经网络相关的函数
import torch.nn as nn
# 提供了一些常用的函数，如softmax
import torch.nn.functional as F
# 优化模块，封装了求解模型的一些优化器，如Adam SGD
import torch.optim as optim
# pytorch 视觉库中提供了一些数据变换的接口
from torchvision import transforms
# pytorch 视觉库提供了加载数据集的接口
from torchvision import datasets
# 用于下载并导入数据集
import torchvision
from torch.utils.data import DataLoader
- 设置超参数
# 由于使用批量训练的方法，需要定义每批的训练的样本数目
batch_size = 64
# 总共训练迭代的次数
epoch = 1
# 设定初始的学习率
LR = 0.001
# 让torch判断是否使用GPU（若使用GPU环境则会更快）
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
- 获取并装载训练集、测试集
# root 用于指定数据集下载后的存放路径
# train 用于指定数据集下载后载入到训练集or测试集（True:载入训练集 False:载入测试集）
# transform 用于指定导入数据集需要对数据进行的操作（转换为tensor数据类型）
# download 设为True表示数据集由程序自动下载
# 下载训练集
train_dataset = datasets.MNIST(root='./data/',
train=True,
transform=transforms.ToTensor(),
download=True)
# 下载测试集
test_dataset = datasets.MNIST(root='./data/',
train=False,
transform=transforms.ToTensor(),
download=True)
# 使用Dataloader数据迭代器加载数据
# dataset 用于指定我们载入的数据集名称
# shuffle 在装载的过程会将数据随机打乱顺序并进行打包，便于后续训练
# 装载训练集
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=True)
- 搭建神经网络
# 卷积神经网络LeNet
# 卷积层使用 torch.nn.Conv2d
# 激活层使用 torch.nn.ReLU
# 池化层使用 torch.nn.MaxPool2d
# 全连接层使用 torch.nn.Linear
class LeNet(nn.Module):
def __init__(self):
super().__init__()
# sequential() 按照顺序排列，要保证相邻层的输入输出大小相匹配
# 卷积 conv2d(in_channels, out_channels, kernel_size, stride, padding)
# in_channels 网络输入的通道数
# out_channels 网络输出的通道数
# kernel_size 卷积核大小q*q
# stride 卷积移动步长
# padding 全零填充
# 激活函数 relu()
# 最大池化 maxpool2d(kernel_size, stride)
# kenel_size 最大池化的窗口大小
# stride 最大池化窗口的移动步长
# 全连接层 linear(in_features, out_features)
# in_features 输入维数
# out_features 输出维数
# BN层 batchnorm1d(num_features)
# num_features 特征维度（等于前一层输出的维度）
self.conv1 = nn.Sequential(nn.Conv2d(1, 6, 3, 1, 2), nn.ReLU(),
nn.MaxPool2d(2, 2))
self.conv2 = nn.Sequential(nn.Conv2d(6, 16, 5), nn.ReLU(),
nn.MaxPool2d(2, 2))
self.fc1 = nn.Sequential(nn.Linear(16 * 5 * 5, 120),
nn.BatchNorm1d(120), nn.ReLU())
self.fc2 = nn.Sequential(
nn.Linear(120, 84),
nn.BatchNorm1d(84),
nn.ReLU(),
# 最后的结果一定要变为 10，因为数字的选项是 0 ~ 9
nn.Linear(84, 10))
# 前向传播
def forward(self, x):
# 卷积+池化
x = self.conv1(x)
# 卷积+池化
x = self.conv2(x)
# 对参数实现扁平化，便于后面全连接层输入
x = x.view(x.size()[0], -1)
# 全连接
x = self.fc1(x)
# 全连接
x = self.fc2(x)
return x
# 将网络操作移动到gpu或者cpu
net = LeNet().to(device)
# 定义交叉熵损失函数
criterion = nn.CrossEntropyLoss()
# 定义模型优化器（Adam自适应优化算法）
# 输入模型参数和初始学习率
optimizer = optim.Adam(
net.parameters(),
lr=LR,
)
- 训练模型
def train(epoch):
sum_loss = 0.0
# 从迭代器抽取图片和标签
for i, data in enumerate(train_loader):
inputs, labels = data
# 将优化器内部参数梯度归零
optimizer.zero_grad()
# 将数据传入网络进行前向运算
outputs = net(inputs)
# 计算损失函数值
loss = criterion(outputs, labels)
# 反向传播
loss.backward()
# 通过梯度做一步参数更新
optimizer.step()
# 计算每批次的损失值
sum_loss += loss.item()
if i % 100 == 99:
print('[%d,%d] loss:%.03f' %
(epoch + 1, i + 1, sum_loss / 100))
sum_loss = 0.0
- 测试模型
def test():
# 将模型变换为测试模式
net.eval()
correct = 0
for data_test in test_loader:
images, labels = data_test
# 测试输入数据得到输出
output_test = net(images)
# 找到概率最大的分类
# torch.max()返回value,index（_表示value,predicted表示index）
# dim=1表示输出所在行（样本）的最大值 dim=0表示输出所在列（类别）的最大值
_, predicted = torch.max(output_test, 1)
# 计算正确分类的数目
correct += (predicted == labels).sum()
# 输出测试集中正确分类的数目
print("correct1: ", correct)
# 输出测试集分类正确率
print("Test acc: {0}".format(correct.item() /
len(test_dataset)))
# 程序入口
if __name__ == '__main__':
for epoch in range(epoch):
train(epoch)
test()
