本文使用pytorch完成一个非常经典的任务——手写数字识别。数据集为LeCun等人于90年代构建的手写数字集MNIST。本文的重点在于对数据的处理，因为torchvision库虽然可以直接加载训练数据，但是这个过程中是看不见数据的，所以本文使用dataset和dataloader库来加载数据，可以从中体会数据处理的过程。
首先引用本项目需要用到的所有库，并定义本项目的一些超参数。
import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader
from torch.optim.lr_scheduler import ExponentialLR
import torch.nn as nn
epochs = 10
batch_size = 128
learning_rate = 0.011 数据处理
1.1 数据集结构
MNIST数据集可以在官方页面下载，整个数据集分为训练/测试图片和训练/测试标签四个文件。这四个文件以二进制的形式存储，对于图片文件来说，它的前16位是图片集的各种信息，后面按顺序存储着图片的像素信息；对于标签文件来说，它的前8位是标签集的各种信息，后面存储着标签值。
1.2 加载训练数据
下载并解压四个文件，即可开始读取数据。根据1.1中的介绍，我们对数据集处理后读入numpy数组。
#数据集路径
images_train = 'D:\\datasets\\MNIST\\train-images.idx3-ubyte'
images_test = 'D:\\datasets\\MNIST\\t10k-images.idx3-ubyte'
labels_train = 'D:\\datasets\\MNIST\\train-labels.idx1-ubyte'
labels_test = 'D:\\datasets\\MNIST\\t10k-labels.idx1-ubyte'
#读取数据到numpy数组
def read_images(dir):
with open(dir, 'rb') as f:
f.read(16)
tmp = f.read()
images = np.frombuffer(tmp, dtype=np.uint8).astype("float32").reshape(int(len(tmp)/(28*28)), 1, 28, 28)
return images
def read_labels(dir):
with open(dir, 'rb') as f:
f.read(8)
tmp = f.read()
labels = np.frombuffer(tmp, dtype=np.uint8).astype("int64")
return labels经过函数read_images的处理，图片集将会被处理为 (C, 1, H, W) 的numpy数组。C表示图片的个数；对于每张图片，它的尺寸是 (1, H, W) ，H、W分别为图片高和宽。将图片集处理成 (C, 1, H, W) 而不是 (C, H, W) 的原因是，将来网络提取的特征是高维的，需要预留一个数组维度。
经过函数read_labels的处理，标签集将会被处理为一维向量。
此时，我们分别读取到了图像和标签信息。接下来，我们需要一一建立图片与标签的对应关系，这里可以用pytorch中的Dataset（torch.utils.data.Dataset）来完成。注意，我们在建立数据集时，首先要继承Dataset类，并且根据官方文档的说明，我们还需要重写__getitem__()和__len__()。可以看到，图像和标签的对应关系在__getitem__()中构建好了。
建立好对应关系后，我们还需要通过DataLoader（torch.utils.data.Dataloader）构建一个迭代器，它可以把整个数据集按照batch_size分成多个batch，在训练时迭代地把这些数据按批次送入网络。
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print("Running on", device, '\n')
class dataset(Dataset):
def __init__(self, images, labels):
self.images = torch.tensor(images).to(device)
self.labels = torch.tensor(labels).to(device)
def __getitem__(self, item):
return self.images[item], self.labels[item]
def __len__(self):
return self.images.shape[0]
train_data = dataset(read_images(images_train), read_labels(labels_train))
train_dataloader = DataLoader(train_data, batch_size=batch_size, shuffle=True, drop_last=False)
test_data = dataset(read_images(images_test), read_labels(labels_test))
test_dataloader = DataLoader(test_data, batch_size=batch_size, shuffle=False, drop_last=False)这段代码前面两行指定了代码运行的设备。这里应该注意，要使用gpu，我们就需要同时把数据和网络权重都储存在gpu上。
2 网络构建
网络部分是该任务的重点，这一部分可以自行设计，例如我自己设计的MyNet：
class MyNet(nn.Module):
def __init__(self):
super(MyNet, self).__init__()
#(1, 28, 28) -> (32, 12, 12)
self.conv1 = nn.Sequential(nn.Conv2d(in_channels=1, out_channels=32, kernel_size=5, stride=1),
nn.BatchNorm2d(32),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2))
#(32, 12, 12) -> (64, 23, 23)
self.convtrans = nn.Sequential(nn.ConvTranspose2d(in_channels=32, out_channels=64, kernel_size=2, stride=4),
nn.BatchNorm2d(64),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2))
#(64, 23, 23) -> (128, 5, 5)
self.conv2 = nn.Sequential(nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, stride=2),
nn.BatchNorm2d(128),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2))
#(128, 5, 5) -> (3200,)
self.flatten = nn.Flatten(start_dim=1, end_dim=-1)
self.linear = nn.Linear(3200, 10)
def forward(self, input):
input = self.conv1(input)
input = self.convtrans(input)
input = self.conv2(input)
input = self.flatten(input)
output = self.linear(input)
return output3 模型加载
这一部分实例化定义好的网络，并将模型加载到运算设备(cpu或gpu)上。同时定义了优化器、学习率衰减策略和损失函数。
model = MyNet()
model.to(device)
optimizer = torch.optim.Adam(model.parameters(), lr = learning_rate)
scheduler = ExponentialLR(optimizer, gamma=0.7)
loss_calculator = nn.CrossEntropyLoss().to(device)4 训练及测试
本文中训练100个batch就执行一次测试，用来观察训练情况。训练及测试代码如下：
def train(train_dataloader, test_dataloader, optimizer, loss_calculator):
for epoch in range(epochs):
for batch, (images, labels) in enumerate(train_dataloader):
preds = model(images)
loss = loss_calculator(preds, labels)
optimizer.zero_grad()
loss.backward()
optimizer.step()
if batch % 100 == 0:
print("Epoch: {: <5}Batch: {: <8}Train_loss: {: <15.6f}".format(epoch, batch, float(loss.data)), end='')
test(test_dataloader)
scheduler.step()
def test(test_dataloader):
match = 0
total = 0
with torch.no_grad():
for imgs, labels in test_dataloader:
preds = model(imgs)
preds = torch.max(preds, 1)[1]
match += int((preds == labels).sum())
total += int(labels.size()[0])
accuracy = match/total
print('Test accuracy: {:.2f}%'.format(accuracy*100))5 运行代码和模型保存
train(train_dataloader, test_dataloader, optimizer, loss_calculator)
torch.save(model, "MyNet_BS{}_LR{}_ExpLR{}.pth".format(batch_size, learning_rate, 0.7))模型的保存和加载有很多方式，本文不再赘述。
本文完整代码在这里：，积分不够的话可以把以上所有代码段复制拼接到一个文档运行。
