训练范式演进

import torch
import torch.nn as nn
import torch.optim as optim

# 1. 数据加载（CPU 内存）
train_dataset = MNIST(root='./data', train=True, download=True)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=32)

# 2. 模型定义（CPU 内存）
model = nn.Sequential(
    nn.Linear(784, 128),
    nn.ReLU(),
    nn.Linear(128, 10)
)
optimizer = optim.SGD(model.parameters(), lr=0.01)

# 3. 将模型移动到 GPU（CPU → GPU 显存）
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = model.to(device)  # 参数从 CPU 内存拷贝到 GPU 显存

for epoch in range(10):
    for batch_images, batch_labels in train_loader:
        # 4. 数据传输（CPU 内存 → GPU 显存，通过 PCIe）
        batch_images = batch_images.to(device)
        batch_labels = batch_labels.to(device)

        # 5. 前向传播（GPU 计算，数据在 GPU 显存）
        logits = model(batch_images.view(-1, 784))
        loss = nn.CrossEntropyLoss()(logits, batch_labels)

        # 6. 反向传播（GPU 计算，梯度在 GPU 显存）
        optimizer.zero_grad()
        loss.backward()  # 自动计算梯度

        # 7. 参数更新（GPU 计算，参数在 GPU 显存）
        optimizer.step()

# 8. 保存模型参数（GPU 显存 → CPU 内存 → 磁盘）
torch.save(model.state_dict(), 'model_weights.pt')
# 内部流程：
# 1. state_dict() 收集所有参数张量（仍在 GPU 显存）
# 2. torch.save() 将张量拷贝到 CPU 内存（通过 PCIe）
# 3. 序列化为字节流并写入磁盘

# 加载模型（磁盘 → CPU 内存 → GPU 显存）
loaded_model = nn.Sequential(
    nn.Linear(784, 128),
    nn.ReLU(),
    nn.Linear(128, 10)
)
loaded_model.load_state_dict(torch.load('model_weights.pt'))  # 磁盘 → CPU 内存
loaded_model.to(device)  # CPU 内存 → GPU 显存

磁盘 → CPU 内存 → PCIe 总线 → GPU 显存
 ↓        ↓          ↓          ↓
读取    预处理     传输(瓶颈)   计算

# 以一个简单的 2 层 MLP 为例
class SimpleMLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(784, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        # 前向传播
        h1 = self.fc1(x)        # 激活值 1：需要保存用于反向传播
        a1 = torch.relu(h1)    # 激活值 2：需要保存用于反向传播
        out = self.fc2(a1)     # 激活值 3：需要保存用于反向传播
        return out

# 反向传播时的依赖链
loss.backward()
# ↓ 需要计算 ∂loss/∂fc2
#   ↓ 需要 a1（前向传播保存的激活值）
# ↓ 需要计算 ∂loss/∂fc1
#   ↓ 需要 h1（前向传播保存的激活值）
#   ↓ 需要 x（输入数据）

总显存 = 模型参数 + 激活值 + 梯度 + 优化器状态

- 模型参数：FP16 存储权重
- 激活值：前向传播的所有中间结果（与 batch size 成正比）
- 梯度：反向传播计算的梯度（与参数量相同）
- 优化器状态：Adam 的动量、方差（通常是参数的 2-3 倍）