背景

这几天在用MMPose训练模型，给ViTPose加一个CNN Stem的时候复杂的配置文件已经要把我整懵逼了，之后要改ViT源码的时候才发现我的pytorch代码水平如此捉急，简直就是啥都要gpt一下，实在惨不忍睹。因此打算开一个pytorch新坑，把官方文档老实过一遍，让自己的pytorch水平不要仅仅纸上谈兵。

文章在Pytorch官方文档的基础上添加了一些知识的进一步讲解，也算部分原创了。

Pytorch机器学习任务工作流

由于这是关于pytorch第一篇文章，因此不上比较困难的部分，走一遍基本的机器学习任务流程就可以。

一个经典机器学习任务分为如下几个步骤：

1. 准备数据
2. 定义模型结构 + 初始化模型
3. 定义损失函数 + 优化器
4. 定义训练和测试过程
5. 开启训练
6. 保存模型
7. 使用保存的模型模型对数据进行推理

下面以FashionMnist作为例子，把这几个步骤从头到尾走一遍。

0. 导包

import torch
from torch import nn # neural network类
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import ToTensor # 转换tensor函数

1. 准备数据

1.1 定义Dataset

先直接上代码：

# 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(),
)

这里使用了torch.utils.data.Dataset这个类。

在抽象层面上，Dataset类存储了样本和它们对应的标签。

但是，从底层存储机制来看，Dataset并没有真正地存储数据本身，而是存储了数据的元信息和访问数据的方法。

大多数情况下，Dataset 对象不会把你的整个数据集（比如几百GB的图像文件）直接加载到内存中。那样会很快耗尽内存。相反，它通常存储的是：

• 文件路径列表： 例如，一个包含所有图片文件路径的 Python 列表。

• 数据库连接信息： 如果数据在数据库中。

• 指向数据块的指针/索引： 如果数据已经被预处理并存储在一个大的二进制文件中(比如ubyte)。

Dataset类的核心在于他的一个私有函数：__getitem__(self, index)，当这个方法被调用时，他会依次进行如下操作：

1. 接受一个整数索引index
2. 根据这个索引，实时地从硬盘、内存、网络或者其他存储介质中加载对应的单个数据样本（一般是从硬盘取）
3. （如果定义了预处理流程），对这个样本进行数据预处理操作，例如归一化，数据增强，裁剪等等。
4. 返回一个处理好的张量Tensor，代表一个独立的输入-标签对。

在底层原理上，Dataset负责的是按需加载，也就是说，只有需要某个特定样本的时候，才会出发对该样本的实际数据读取和处理操作。这使得电脑可以处理比内存大得多的数据集。

1.2 定义Dataloader

刚才的Dataset定义了数据在哪里，以及一个根据索引来取数据样本的接口。但是实际你要怎么把数据按需搬到内存里面呢？这就得用到Dataloader了。

batch_size = 64
num_workers=10

# Create data loaders.
train_dataloader = DataLoader(training_data, batch_size=batch_size, shuffle=True, num_workers=7, pin_memory=False, persistent_workers=True)
test_dataloader = DataLoader(test_data, batch_size=batch_size, shuffle=True, num_workers=7, pin_memory=False, persistent_workers=True)

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

Dataloader关注的是数据从 Dataset 到模型的高效流动。它利用了操作系统的内存管理和进程/线程调度能力。

在一次批处理中，Dataloader会循环调用Dataset的__getitem__方法，次数等于batch_size，把一次批处理要处理的样本数量全部搬到计算内存中，并且把这些独立的样本收集起来，拼接成一个更大的张量， 这个张量就是一个batch，是训练时数据的最小单元。例如，如果 __getitem__ 返回一个 (C, H, W) 的图像张量，那么 DataLoader 会将其拼接成 (batch_size, C, H, W)。这个拼接过程主要是发生在CPU内存中。

Dataloader有4个关键操作：

1. 多进程加载：也就是那个num_workers参数的设置，这是加速训练你能做的最重要的优化之一， 当设置num_workers > 1的时候，Dataloader会启动num_workers数量的子进程，每个子进程都拥有一个独立的Dataset实例。这些子进程会并行地执行 __getitem__ 方法，从硬盘加载数据、进行预处理，并将处理好的单个样本传回主进程。这样，如果num_worker=1，单位时间内你可能就只能取一个样本，但是如果num_worker=10，那单位时间就可以取10个样本，大大减少了数据I/O所产生的数据延迟，加速训练过程。
2. 设置锁页内存
   可以通过pin_memory=True来启用锁页内存, 它会指示 PyTorch 将数据加载到锁页内存（GPU可以直接访问，而不是要从CPU侧拷贝到GPU侧才能访问），从而加快数据从 CPU 到 GPU 的传输速度，减少 GPU 等待数据的时间。
3. 设置持久workers:
   可以指定persistent_worker=True进一步加速训练，其原理为：默认情况下，每个 epoch 结束后，DataLoader 的子进程都会被销毁并重新创建。这个过程会有一定的开销。设置 persistent_workers=True 会让这些子进程在整个训练过程中保持活动状态，避免了重复创建和销毁的开销，这对于短 epoch 或大量 epoch 的训练场景特别有用。
4. 数据打乱
   DataLoader 内部会维护一个索引列表。当 shuffle=True 时，在每个 epoch 开始时，这个索引列表会被随机打乱。然后，DataLoader 会按照打乱后的索引顺序去请求 Dataset 的 __getitem__ 方法。这样确保了每个 epoch 的数据输入顺序都是不同的，有助于模型更好地泛化。

2. 定义+实例化模型

在Pytorch中，想要定义一个神经网络，或者一个子模块，你就需要创建一个从nn.Module类继承过来的类。不过在这之前，你要先定义你的模型到底会在哪里跑，开头的代码会自动检测你有没有加速芯片，没有的话就会自动用GPU：

device = torch.accelerator.current_accelerator().type if torch.accelerator.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)

一般我们是在类的__init__方法中定义网络有哪些模块/层结构， 需要注意的是这里并没有定义模块/层之间的先后顺序。这里是定义了一个展平层nn.Flatten模块，又定义了一个顺序容器模块nn.Sequential，这个Sequential模块要注意一下，所有被包含在里面的模块都会被一次性顺序执行。

定义完有什么模块之后，下一步就是定义数据在模块间的流动顺序了。这个过程由类中的forward方法实现。

最后就是把自定义的类给实例化并且放到device侧，然后就可以把model信息给打印出来了：

3. 定义损失函数/优化器

要训练一个模型，你得告诉模型他现在的表现是好是坏，以及他应该怎么改。这两个功能对应的实现就是 loss_fn损失函数的定义，以及optimizer优化器的定义。

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

由于是一个分类任何，所以直接用交叉熵损失即可。optimizer直接用随机梯度下降就可以了，学习率lr设置为1e-3。

4. 定义训练/测试过程

4.1 定义训练过程

定义一个训练过程，要指定如下内容：

• 用什么dataloader
• 用什么model
• 用什么loss_fn
• 用什么optimizer

上面的一般作为参数传入过程函数，在函数内部要手动实现

• dataloader对数据集的遍历
• 前向传播
• 损失函数的计算
• 反向传播的梯度计算
• 优化器对模型权重的步进
• 梯度清零

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}]")

4.2 定义测试过程

定义一个测试过程，要指定如下内容：

• 用什么dataloader
• 用什么model
• 用什么loss_fn
  相对于训练过程，测试过程由于只是推理，所以无需使用optimizer更新参数。

上面的一般作为参数传入过程函数，在函数内部要手动实现

• dataloader对数据集的遍历
• 前向传播
• 损失函数的计算
• 正确率的计算

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

5. 开启训练过程

• 指定epoch，在每个epoch进行train过程和test过程

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

6. 保存模型

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

7. 加载并利用模型进行推理

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

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}"')