0基础学习PyTorch——GPU上训练和推理

在《Windows Subsystem for Linux——支持cuda能力》一文中，我们让开发环境支持cuda能力。现在我们要基于《0基础学习PyTorch——时尚分类（Fashion MNIST）训练和推理》，将代码修改成支持cuda的训练和推理。

创建设备

我们首先需要依据环境是否支持cuda来创建相应设备。

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

因为我们开发环境WSL已经支持了cuda，所以此时我们创建的是GPU设备。

训练

训练的过程有两处修改：

• 将模型实例化到GPU上。

model = GarmentClassifier().to(device) # model = GarmentClassifier()

• 将数据移动到GPU上。

inputs, labels = data  # 获取输入数据和对应的标签
inputs, labels = inputs.to(device), labels.to(device)  # 将数据移动到GPU上

完整代码如下

from datetime import datetime
import torch
import torchvision
import torchvision.transforms as transforms
from garmentclassifier import GarmentClassifier

# 定义图像转换操作：将图像转换为张量，并进行归一化处理
transform = transforms.Compose(
    [transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))]) # 对图像的每个通道进行标准化，使得每个通道的像素值具有零均值和单位标准差

# 加载FashionMNIST训练数据集，并应用定义的图像转换操作
training_set = torchvision.datasets.FashionMNIST('./data', train=True, transform=transform)

# 创建数据加载器，用于批量加载训练数据，batch_size为4，数据顺序随机打乱
trainloader = torch.utils.data.DataLoader(training_set, batch_size=4, shuffle=True)

# 将模型移动到GPU上
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 实例化模型并移动到GPU上
model = GarmentClassifier().to(device)

# 定义损失函数为交叉熵损失
loss_fn = torch.nn.CrossEntropyLoss()
# 定义优化器为随机梯度下降（SGD），学习率为0.001，动量为0.9
optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

# 训练模型，训练2个epoch
for epoch in range(2):
    running_loss = 0.0  # 初始化累计损失
    # 枚举数据加载器中的数据，i是批次索引，data是当前批次的数据
    for i, data in enumerate(trainloader, 0):
        inputs, labels = data  # 获取输入数据和对应的标签
        
        inputs, labels = inputs.to(device), labels.to(device)  # 将数据移动到GPU上
        
        optimizer.zero_grad()  # 清空梯度
        outputs = model(inputs)  # 前向传播，计算模型输出
        loss = loss_fn(outputs, labels)  # 计算损失
        loss.backward()  # 反向传播，计算梯度
        optimizer.step()  # 更新模型参数
        running_loss += loss.item()  # 累加损失

        # 每2000个批次打印一次平均损失
        if i % 2000 == 1999:
            print(f'[{epoch + 1}, {i + 1}] loss: {running_loss / 2000}')
            running_loss = 0.0  # 重置累计损失
      
# 获取当前时间戳，格式为 'YYYYMMDD_HHMMSS'
timestamp = datetime.now().strftime('%Y%m%d%H%M%S.pth')

# 定义模型保存路径，包含时间戳
model_path = 'model_{}'.format(timestamp)      

# 保存模型的状态字典到指定路径
torch.save(model.state_dict(), model_path)

推理

GPU上算出的模型不一定非要在GPU上推理，也可以在CPU上推理。
但是本文我们就是希望模型在GPU上推理，则可以对代码做如下修改。

• 将模型实例化到GPU上。

model = GarmentClassifier().to(device)  # model = GarmentClassifier()

• 将数据移动到GPU上。

image = image.to(device)  # 将图像移动到GPU上

完整代码如下

import os
import glob
import torch
import torchvision.transforms as transforms
from PIL import Image
from datetime import datetime
from garmentclassifier import GarmentClassifier

def get_latest_model_path(directory, pattern="model_*.pth"):
    # 获取目录下所有符合模式的文件
    model_files = glob.glob(os.path.join(directory, pattern))
    if not model_files:
        raise FileNotFoundError("No model files found in the directory.")
    
    # 找到最新的模型文件
    latest_model_file = max(model_files, key=os.path.getmtime)
    return latest_model_file

# 定义图像转换操作：将图像转换为张量，并进行归一化处理
transform = transforms.Compose([
    transforms.Resize((28, 28)),  # 调整图像大小为28x28
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

# 将模型移动到GPU上
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 实例化模型并移动到GPU上
model = GarmentClassifier().to(device)  

# 加载训练好的模型
model_path = get_latest_model_path('./')  # 获取最新的模型文件
model.load_state_dict(torch.load(model_path, weights_only=False)) # 加载模型参数
model.eval()  # 设置模型为评估模式

# 从本地加载图像
image_path = 'shoe.jpg'  # 替换为实际的图像路径
image = Image.open(image_path).convert('L')  # 将图像转换为灰度图

# 预处理图像
image = transform(image)
image = image.unsqueeze(0)  # 增加一个批次维度
image = image.to(device)  # 将图像移动到GPU上

# 推理（预测）
with torch.no_grad():  # 在推理过程中不需要计算梯度
    outputs = model(image)  # 前向传播，计算模型输出
    _, predicted = torch.max(outputs, 1)  # 获取预测结果

# 定义类别名称
classes = ('T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat',
           'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot')

# 打印预测结果
print(f'Predicted label: {classes[predicted.item()]}')

总结

• 依据系统是否支持cuda来生成设备。
• 模型和数据都要移动到相同的设备上。
• 模型是由CPU还是GPU训练的，并不影响推理使用CPU还是GPU。