用PyTorch完成手写数字识别

输出结果：

Epoch: [2/3] Iteration: [850/938] loss: 0.023803 Epoch: [2/3] Iteration: [900/938] loss: 0.021752 第2个epoch的识别准确率为：99.4000% Epoch: [3/3] Iteration: [50/938] loss: 0.013843 Epoch: [3/3] Iteration: [100/938] loss: 0.014237 Epoch: [3/3] Iteration: [150/938] loss: 0.016309

Loss 可视化：

代码：

import torch import torch.nn as nn import torchvision as tv from torch.utils import data from torchvision import transforms as T import os import numpy as np from PIL import Image import matplotlib.pyplot as plt import tensorboardX from utils import show_confMat, validate device = torch.device("cuda" if torch.cuda.is_available() else "cpu") transform = T.Compose([T.Resize(224), T.ToTensor()]) # 获取训练数据集 data_train = tv.datasets.MNIST( root='./data/mnist/', train=True, download=True, transform=transform) trainloader = data.DataLoader( data_train, batch_size=64, shuffle=True, ) # 获取测试数据集 data_test = tv.datasets.MNIST( root='./data/mnist/', train=False, download=True, transform=transform) testloader = data.DataLoader( data_test, batch_size=64, shuffle=False, ) # 获取预选连好的resnet18,修改第一层卷积层和最后一层全连接层 net = tv.models.resnet18(True) net.conv1 = nn.Conv2d(1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False) net.fc=nn.Linear(512, 10) net.to(device) # 交叉熵损失函数，通常用于多分类问题上 import torch.optim as optim criterion = nn.CrossEntropyLoss() optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9) # tensorboard可视化 from tensorboardX import SummaryWriter writer = SummaryWriter(logdir="./log/") # 训练 + 测试 for epoch in range(3): sum_loss = 0.0 epoch_loss = 0.0 for i, data in enumerate(trainloader): img, labeles = data img, labeles = img.to(device), labeles.to(device) # 梯度清零 optimizer.zero_grad() # forward + backward outputs = net(img) loss = criterion(outputs, labeles) loss.backward() optimizer.step() # 每训练50个batch打印一次平均loss sum_loss += loss.item() epoch_loss += loss.item() if (i+1) % 50 == 0: print('Epoch: [{}/{}] Iteration: [{}/{}] loss: {:6f}'.\ format(epoch+1, 3, i+1, len(trainloader), sum_loss / 50)) writer.add_scalars('Loss', {'train_loss': sum_loss / 50}, epoch*len(trainloader)+i+1) sum_loss = 0.0 # 每跑完一次epoch测试一下准确率 with torch.no_grad(): correct = 0 total = 0 for data in testloader: images, labels = data images, labels = images.to(device), labels.to(device) outputs = net(images) # 取得分最高的那个类 _, predicted = torch.max(outputs.data, 1) total += labels.size(0) correct += (predicted == labels).sum().item() print('第{}个epoch的识别准确率为：{:.4%}'.format(epoch+1, 100.0 * correct / total)) torch.save(net.state_dict(), './net_{}.pth'.format(epoch+1))