本章利用前述PyTorch的基本使用方法,来完成一个对MNIST数据集的手写数字图片分类任务。
首先我们继承torch.nn.Module
类,创建一个自定义的网络模型,继承时至少需要重写两个方法:__init__()
和forward(x)
,前者用于模型初始化,定义模型结构;后者定义前向传播的计算步骤,即输入x如何得到输出y。我们的实现如下:
class Model(nn.Module):
def __init__(self):
# 调用父类的方法
super(Model, self).__init__()
# 定义模型结构
self.layers = nn.Sequential(
nn.Flatten(), # 将训练样本输入转化为一维向量
nn.Linear(28*28, 128),
nn.ReLU(),
nn.Linear(128, 64),
nn.ReLU(),
nn.Linear(64, 10),
nn.ReLU(),
)
def forward(self, x):
return self.layers(x)
上述模型结构的隐藏层的层数和神经元数量、激活函数等,都可以自由设置,经过实验选择分类准确率最高的配置。
下面定义每个epoch所需要做的工作。一共分为4步:前向传播得到输出、根据输出计算损失、反向传播计算导数、根据导数更新模型参数。在实际的实现中,为了消除上一个epoch所计算的导数对当前epoch计算倒数的影响,需要清空之前的epoch累计的梯度。实现如下:
def train_loop(dataloader: DataLoader, model: nn.Module, loss_fn, optimizer: torch.optim.Optimizer):
size = len(dataloader.dataset)
for batch, (x, y) in enumerate(dataloader):
# 前向传播
y1 = model(x)
# 计算损失
loss = loss_fn(y1, y)
# 清空累计梯度
optimizer.zero_grad()
# 反向传播
loss.backward()
# 更新参数
optimizer.step()
# 输出结果
if batch % 200 == 0:
loss, current = loss.item(), batch * len(x)
print(f"loss: {loss:>7f} [{current:>5d}/{size:>5d}]")
模型测试时的步骤与训练相似,只是不用计算梯度、更新参数,于是在这一步可以禁用自动求导,使得运算更加迅速。实现如下:
def test_loop(dataloader: DataLoader, model: nn.Module, loss_fn, optimizer: torch.optim.Optimizer):
# 初始化损失函数值、分类准确的样本数
test_loss, correct = 0, 0
# 获取dataloader的batch总数
num_batches = len(dataloader)
# 获取测试数据样本总数
size = len(dataloader.dataset)
# 禁用梯度,提高计算速度
with torch.no_grad():
# 遍历dataloader中的所有数据,注意x,y依然是一组batch而不是单个样本
for x, y in dataloader:
# 前向传播
y1 = model(x)
# 计算损失
test_loss += loss_fn(y1, y).item()
# 计算预测正确的数量,对一组batch的分类正确数量求和
correct += (y1.argmax(1) == y).type(torch.float).sum().item()
test_loss /= num_batches
correct /= size
print(f"Test Error: \n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f} \n")
加载训练集和测试集的数据,注意将PIL转换为Tensor:
# 读取训练数据
train_data = datasets.MNIST(
root="./datas/",
download=False,
train=True,
transform=transforms.ToTensor()
)
# 读取测试数据
test_data = datasets.MNIST(
root="./datas/",
download=False,
train=False,
transform=transforms.ToTensor()
)
# 装载训练和测试数据,并设置batch_size,每次读取乱序
train_loader = DataLoader(train_data, shuffle=True, batch_size=64)
test_loader = DataLoader(test_data, shuffle=True, batch_size=64)
初始化模型实例,并设置运行的设备,如果计算机上安装了gpu版本的PyTorch,将会使用GPU加速运算:
# 设置要运行的设备,若安装了cuda则使用gpu,否则使用cpu
device = 'cuda' if torch.cuda.is_available() else 'cpu'
# 初始化模型,并设置运行的设备
model = Model().to(device)
初始化优化器对象,这里使用SGD,除此之外还有Adam、Adagrad等可以使用。
# 初始化优化器,使用SGD优化器
optimizer = torch.optim.SGD(model.parameters(), lr=lr)
损失函数使用交叉熵损失函数,注意使用交叉熵损失函数时,数据集中的label不必映射为one-hot编码,保持原样使用一个整型值表示第几类即可。
# 设置损失函数,使用交叉熵
loss_fn = nn.CrossEntropyLoss()
训练模型只需要使用一个for循环,按设定好的迭代次数执行即可(下面仅显示最后训练完成输出的结果):
# 开始训练
for i in range(epoch):
print(f"Epoch {i+1}\n-------------------------------")
train_loop(train_loader, model, loss_fn, optimizer)
test_loop(test_loader, model, loss_fn, optimizer)
...
Epoch 5
-------------------------------
Accuracy: 97.2%, Avg loss: 0.089015
完整代码请参见MNIST分类-CSDN完整代码资源。