1. 加载数据集和数据预处理
首先导入需要的包和第三方库;
import torch
import torch.nn as nn
import numpy as np
import torchvision
from torch.utils.data import DataLoader
然后设置分批次训练数据的大小、训练次数和学习率;
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
# 分批次训练,一批 64 个训练数据
BATCH_SIZE = 64
# 所有训练数据训练 3 次
EPOCHS = 50
# 学习率设置为 0.0001
LEARN_RATE = 1e-4
加载数据集,将下载的文件转换成pytorch认识的tensor类型,并将图片的数值大小从0—255归一化到0—1;
# 加载数据集
train_data = torchvision.datasets.MNIST(
root='./mnist',
train=True,
transform=torchvision.transforms.ToTensor() ,# 将下载的文件转换成pytorch认识的tensor类型,且将图片的数值大小从(0-255)归一化到(0-1)
download=True
)
train_loader = DataLoader(dataset=train_data, batch_size=BATCH_SIZE, shuffle=True)
test_data = torchvision.datasets.MNIST(
root='./mnist',
train=False,
transform=torchvision.transforms.ToTensor() # 将下载的文件转换成pytorch认识的tensor类型,且将图片的数值大小从(0-255)归一化到(0-1)
)
test_loader = DataLoader(dataset=test_data, batch_size=1, shuffle=True)
2. 定义CNN网络结构
定义CNN神经网络结构,定义了两个卷积层、一个全连接输出层,第一个卷积层输入通道为1,输出通道为16,卷积核大小为5,使用Relu激活函数和最大值池化。第二个卷积层输入通道为16,输出通道为32,卷积核大小为5,使用Relu激活函数和最大值池化。全连接层,输入大小为32*7*7,输出大小为10,并定义向前传播的方法。
class CNN(nn.Module):
def __init__(self):
super(CNN, self).__init__()
# 卷积层1:输入通道为 1,输出通道为 16,卷积核大小 为 5
# 使用 Relu 激活函数
# 使用最大值池化
self.conv1 = nn.Sequential(
nn.Conv2d(
in_channels=1,
out_channels=16,
kernel_size=5,
stride=1,
padding=2
),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2)
)
# 卷积层2:输入通道为 16,输出通道为 32,卷积核大小 为 5
# 使用 Relu 激活函数
# 使用最大值池化
self.conv2 = nn.Sequential(
nn.Conv2d(
in_channels=16,
out_channels=32,
kernel_size=5,
stride=1,
padding=2
),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2),
)
# 输出层,全连接层,输入大小 32 * 7 * 7, 输出大小 10
self.layer_out = nn.Linear(32 * 7 * 7, 10)
def forward(self, x):
x = self.conv1(x)
x = self.conv2(x)
x = x.view(x.size(0), -1)
self.out = self.layer_out(x)
return self.out
使用CNN模型在GPU上进行训练,使用交叉熵损失函数,并使用Adam优化器进行优化。
# 实例化CNN,并将模型放在 GPU 上训练
model = CNN().to(device)
# 使用交叉熵损失,同样,将损失函数放在 GPU 上
loss_fn = nn.CrossEntropyLoss().to(device)
# 使用 Adam 优化器
optimizer = torch.optim.Adam(model.parameters(), lr=LEARN_RATE)
加载数据集并调用模型进行训练,计算并输出损失值。在每一次循环之前,将梯度清零,进行反向传播调整权重,梯度下降来求解损失函数的最小值。此处设置batch_size为64,所以每次训练共有938批次。
for epoch in range(EPOCHS):
# 加载训练数据
for step, data in enumerate(train_loader):
x, y = data
x, y = x.to(device), y.to(device)
# 调用模型预测
output = model(x).to(device)
# 计算损失值
loss = loss_fn(output, y.long())
# 输出看一下损失变化
print(f'EPOCH({epoch}) step({step}) loss = {loss.item()}')
# 每一次循环之前,将梯度清零
optimizer.zero_grad()
# 反向传播
loss.backward()
# 梯度下降
optimizer.step()
使用模型测试数据,对10000张测试集图像进行预测输出,得到10个输出值,即为该图片为每个数字的概率,取最大概率的作为预测值,将其与真实值进行比较,预测准确则sum+1,最后计算并输出预测准确率。
sum = 0
# test:
for i, data in enumerate(test_loader):
x, y = data
x, y = x.to(device), y.to(device)
# 得到模型预测输出,10个输出,即该图片为每个数字的概率
res = model(x)
# 最大概率的就为预测值
r = torch.argmax(res)
l = y.item()
sum += 1 if r == l else 0
print(f'test({i}) CNN:{r} -- label:{l}')
print('accuracy:', sum / 10000)
训练模型的输出结果如下:
对测试集进行测试的输出结果如下: