使用 PyTorch 构建机器学习应用

通过 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)

这样模型就创建好了。执行上面代码,可以看到输出模型的结构信息:

NeuralNetwork(
  (flatten): Flatten(start_dim=1, end_dim=-1)
  (linear_relu_stack): Sequential(
    (0): Linear(in_features=784, out_features=512, bias=True)
    (1): ReLU()
    (2): Linear(in_features=512, out_features=512, bias=True)
    (3): ReLU()
    (4): Linear(in_features=512, out_features=10, bias=True)
  )
)

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!")

执行代码输出信息,如下所示:

Epoch 1
-------------------------------
loss: 2.290439  [   64/60000]
loss: 2.286898  [ 6464/60000]
loss: 2.269897  [12864/60000]
loss: 2.277002  [19264/60000]
loss: 2.252169  [25664/60000]
loss: 2.215247  [32064/60000]
loss: 2.227317  [38464/60000]
loss: 2.184588  [44864/60000]
loss: 2.197558  [51264/60000]
loss: 2.157371  [57664/60000]
Test Error: 
 Accuracy: 35.2%, Avg loss: 2.155970 

Epoch 2
-------------------------------
loss: 2.160764  [   64/60000]
loss: 2.154447  [ 6464/60000]
loss: 2.099050  [12864/60000]
loss: 2.120136  [19264/60000]
loss: 2.069308  [25664/60000]
loss: 2.003359  [32064/60000]
loss: 2.033523  [38464/60000]
loss: 1.949132  [44864/60000]
loss: 1.970084  [51264/60000]
loss: 1.883966  [57664/60000]
Test Error: 
 Accuracy: 56.9%, Avg loss: 1.888517 

Epoch 3
-------------------------------
loss: 1.921510  [   64/60000]
loss: 1.886257  [ 6464/60000]
loss: 1.776533  [12864/60000]
loss: 1.817366  [19264/60000]
loss: 1.713696  [25664/60000]
loss: 1.663400  [32064/60000]
loss: 1.689708  [38464/60000]
loss: 1.586672  [44864/60000]
loss: 1.622822  [51264/60000]
loss: 1.503238  [57664/60000]
Test Error: 
 Accuracy: 61.4%, Avg loss: 1.527024 

Epoch 4
-------------------------------
loss: 1.595791  [   64/60000]
loss: 1.551552  [ 6464/60000]
loss: 1.411490  [12864/60000]
loss: 1.479247  [19264/60000]
loss: 1.368348  [25664/60000]
loss: 1.363437  [32064/60000]
loss: 1.379028  [38464/60000]
loss: 1.298237  [44864/60000]
loss: 1.341409  [51264/60000]
loss: 1.227410  [57664/60000]
Test Error: 
 Accuracy: 63.6%, Avg loss: 1.259987 

Epoch 5
-------------------------------
loss: 1.338061  [   64/60000]
loss: 1.311606  [ 6464/60000]
loss: 1.154888  [12864/60000]
loss: 1.256057  [19264/60000]
loss: 1.137579  [25664/60000]
loss: 1.162467  [32064/60000]
loss: 1.184271  [38464/60000]
loss: 1.114591  [44864/60000]
loss: 1.162392  [51264/60000]
loss: 1.063866  [57664/60000]
Test Error: 
 Accuracy: 64.9%, Avg loss: 1.092105 

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"

参考资源

Creative Commons License

本文基于署名-非商业性使用-相同方式共享 4.0许可协议发布,欢迎转载、使用、重新发布,但务必保留文章署名时延军(包含链接:http://shiyanjun.cn),不得用于商业目的,基于本文修改后的作品务必以相同的许可发布。如有任何疑问,请与我联系

发表评论

电子邮件地址不会被公开。 必填项已用*标注

您可以使用这些HTML标签和属性: <a href="" title=""> <abbr title=""> <acronym title=""> <b> <blockquote cite=""> <cite> <code> <del datetime=""> <em> <i> <q cite=""> <strike> <strong>