IT博客汇 | 使用 PyTorch 构建机器学习应用

使用 PyTorch 构建机器学习应用

Yanjun发表于 2024-01-18 10:55:43

通过 PyTorch 官网给出的 Quickstart，了解使用 PyTorch 完成一个模型的数据准备、模型训练和评估、模型加载并应用。在实际应用中，只需要按照这个流程来编程构建即可。
下面，我们通过分步骤来说明开发机器学习应用程序的基本流程。

我们使用 PyTorch-2.0.1，安装该版本 PyTorch 执行如下命令：

pip install torch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2

1 准备数据

PyTorch 使用 torch.utils.data.DataLoader 和 torch.utils.data.Dataset 来实现数据的加载，并转换成包含样本和标签的 DataSet。
我们可以在 https://pytorch.org/vision/stable/datasets.html 中找到 Torchvision 提供的大量内置 DataSet，通过这些工具类就可以方便构建并使用 DataSet。如果是使用我们自己的数据集，可以继承自 torch.utils.data.Dataset 实现我们自己的 DataSet 以及 DataLoader。
这里使用了 FashionMNIST 数据集，这个数据集是一个包含时尚衣物图像，其中包括衣物图像和它们对应的标签.
首先，下载 FashionMNIST 数据集，可以直接通过 PyTorch 的 datasets API 下载，代码如下：

import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import ToTensor

# Download training data from open datasets.
training_data = datasets.FashionMNIST(
    root="data",
    train=True,
    download=True,
    transform=ToTensor(),
)

# Download test data from open datasets.
test_data = datasets.FashionMNIST(
    root="data",
    train=False,
    download=True,
    transform=ToTensor(),
)

执行上述代码，数据集会自动下载到本地的 data 目录下。
然后我们就可以基于上面得到的 DataSet 创建 DataLoader，代码如下：

batch_size = 64

# Create data loaders.
train_dataloader = DataLoader(training_data, batch_size=batch_size)
test_dataloader = DataLoader(test_data, batch_size=batch_size)

for X, y in test_dataloader:
    print(f"Shape of X [N, C, H, W]: {X.shape}")
    print(f"Shape of y: {y.shape} {y.dtype}")
    break

可以看到，测试集中数据的 shape 信息，一个 batch 有 64 个样本数据，每个样本数据是 1x28x28 的张量：

Shape of X [N, C, H, W]: torch.Size([64, 1, 28, 28])
Shape of y: torch.Size([64]) torch.int64

我们可以通过下面代码查看，训练集中的样例图像和标签数据：

import matplotlib.pyplot as plt

labels_map = {
    0: "T-Shirt",
    1: "Trouser",
    2: "Pullover",
    3: "Dress",
    4: "Coat",
    5: "Sandal",
    6: "Shirt",
    7: "Sneaker",
    8: "Bag",
    9: "Ankle Boot",
}
figure = plt.figure(figsize=(8, 8))
cols, rows = 3, 3
for i in range(1, cols * rows + 1):
    sample_idx = torch.randint(len(training_data), size=(1,)).item()
    img, label = training_data[sample_idx]
    figure.add_subplot(rows, cols, i)
    plt.title(labels_map[label])
    plt.axis("off")
    plt.imshow(img.squeeze(), cmap="gray")
plt.show()

图像数据样例，如下所示：
FashionMNIST

2 创建模型

在 PyTorch 中创建神经网络模型，需要创建一个 Python 类继承自 nn.Module，其中需要在 __init__ 函数中定义神经网络的各个层，并通过实现 forward 函数来指定数据如何传递给神经网络。
实现的神经网络 NeuralNetwork 类，代码如下：

# Get cpu, gpu or mps device for training.
device = (
    "cuda"
    if torch.cuda.is_available()
    else "mps"
    if torch.backends.mps.is_available()
    else "cpu"
)
print(f"Using {device} device")

# Define model
class NeuralNetwork(nn.Module):
    def __init__(self):
        super().__init__()
        self.flatten = nn.Flatten()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(28*28, 512),
            nn.ReLU(),
            nn.Linear(512, 512),
            nn.ReLU(),
            nn.Linear(512, 10)
        )

    def forward(self, x):
        x = self.flatten(x)
        logits = self.linear_relu_stack(x)
        return logits

model = NeuralNetwork().to(device)
print(model)

这样模型就创建好了。
PyTorch 在 torch.nn 模块下面提供了方便构建各种神经网络模型的 building block，使用起来非常灵活：

Containers
Convolution Layers
Pooling layers
Padding Layers
Non-linear Activations (weighted sum, nonlinearity)
Non-linear Activations (other)
Normalization Layers
Recurrent Layers
Transformer Layers
Linear Layers
Dropout Layers
Sparse Layers
Distance Functions
Loss Functions
Vision Layers
Shuffle Layers
DataParallel Layers (multi-GPU, distributed)
Utilities
Quantized Functions
Lazy Modules Initialization

3 训练模型

在训练模型之前，我们需要顶一个 loss function 和一个 optimizer，示例代码如下：

loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)

然后，就可以实现训练模型的处理逻辑。
训练模型的代码，如下所示：

def train(dataloader, model, loss_fn, optimizer):
    size = len(dataloader.dataset)
    model.train()
    for batch, (X, y) in enumerate(dataloader):
        X, y = X.to(device), y.to(device)

        # Compute prediction error
        pred = model(X)
        loss = loss_fn(pred, y)

        # Backpropagation
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

        if batch % 100 == 0:
            loss, current = loss.item(), (batch + 1) * len(X)
            print(f"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]")

使用测试集来评估，上面我们训练得到的模型的性能，代码如下：

def test(dataloader, model, loss_fn):
    size = len(dataloader.dataset)
    num_batches = len(dataloader)
    model.eval()
    test_loss, correct = 0, 0
    with torch.no_grad():
        for X, y in dataloader:
            X, y = X.to(device), y.to(device)
            pred = model(X)
            test_loss += loss_fn(pred, y).item()
            correct += (pred.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")

接下来，我们使用上面定义的训练、评估函数，通过多轮迭代让模型不断地去学习参数，逐步提高模型精度，即最小化 loss 函数的值：

epochs = 5
for t in range(epochs):
    print(f"Epoch {t+1}\n-------------------------------")
    train(train_dataloader, model, loss_fn, optimizer)
    test(test_dataloader, model, loss_fn)
print("Done!")

4 保存模型

通过训练调优，我们得到了我们需要的模型，这时应该将模型保存下来，供后面的应用场景来使用该模型。
通过下面代码来保存模型：

torch.save(model.state_dict(), "model.pth")
print("Saved PyTorch Model State to model.pth")

保存成功以后，可以在程序运行的当前目录下，看到名称为 model.pth 的模型文件，这个可以在后面加载模型的时候直接使用。
在实际应用中，我们可能会随着数据的积累，能够不断地训练调优模型，从而又会生成新的精度更高更实用的模型，这时在保存模型的时候可以设置模型的版本号，用来标识模型的升级，然后就可以在使用模型的时候替换旧的模型。

5 加载模型

加载模型，代码如下：

model = NeuralNetwork().to(device)
model.load_state_dict(torch.load("model.pth"))

模型加载成功以后，我们可以使用这个 model 来进行预测（Make Predictions）：

classes = [
    "T-shirt/top",
    "Trouser",
    "Pullover",
    "Dress",
    "Coat",
    "Sandal",
    "Shirt",
    "Sneaker",
    "Bag",
    "Ankle boot",
]

model.eval()
x, y = test_data[0][0], test_data[0][1]
with torch.no_grad():
    x = x.to(device)
    pred = model(x)
    predicted, actual = classes[pred[0].argmax(0)], classes[y]
    print(f'Predicted: "{predicted}", Actual: "{actual}"')

上面代码中，使用了测试集中一个图像数据 test_data[0][0] 来进行示例预测，实际中可以输入其它我们想要预测的图像数据。
执行代码，输出预测结果：

Predicted: "Ankle boot", Actual: "Ankle boot"

参考资源