python学习体系

Appearance

章节2：PyTorch深度学习框架 ​

学习目标 ​

  • 掌握 PyTorch 张量操作与自动微分机制
  • 学会使用 Dataset 和 DataLoader 构建自定义数据集
  • 理解 nn.Module 模型定义与完整训练循环
  • 能够利用 GPU 加速模型训练
  • 掌握模型保存/加载与 ONNX 导出

2.1 张量操作与自动微分 ​

PyTorch 的核心是 张量（Tensor），类似 NumPy 的 ndarray，但支持 GPU 加速和自动微分。

2.1.1 创建张量 ​

python
import torch

# 从列表创建
a = torch.tensor([[1, 2], [3, 4]])
print(a)

# 特殊张量
zeros = torch.zeros(3, 4)        # 全零
ones = torch.ones(2, 3)          # 全一
eye = torch.eye(5)               # 单位矩阵
rand = torch.randn(3, 3)         # 标准正态分布

# 形状操作
x = torch.randn(2, 3, 4)
print(x.shape)                   # torch.Size([2, 3, 4])
print(x.view(2, -1).shape)       # torch.Size([2, 12]) 展平

2.1.2 张量运算 ​

python
x = torch.tensor([1.0, 2.0, 3.0])
y = torch.tensor([4.0, 5.0, 6.0])

print(x + y)       # 加法
print(x * y)       # 逐元素乘
print(x @ y)       # 点积（矩阵乘法用 @ 或 torch.mm）
print(torch.cat([x, y], dim=0))  # 拼接

# 广播机制
a = torch.randn(3, 1)
b = torch.randn(1, 4)
print((a + b).shape)  # torch.Size([3, 4]) — 广播到相同形状

2.1.3 自动微分（Autograd） ​

PyTorch 的 autograd 机制自动记录计算图并计算梯度。

python
x = torch.tensor([2.0], requires_grad=True)
w = torch.tensor([3.0], requires_grad=True)
b = torch.tensor([1.0], requires_grad=True)

# 前向计算
y = w * x + b
loss = y.mean()

# 反向传播
loss.backward()

# 查看梯度
print(f"dx: {x.grad}")    # dy/dx = w = 3.0
print(f"dw: {w.grad}")    # dy/dw = x = 2.0
print(f"db: {b.grad}")    # dy/db = 1.0

# 关闭梯度追踪（推理时）
with torch.no_grad():
    y_pred = w * x + b

2.2 Dataset 与 DataLoader ​

2.2.1 自定义 Dataset ​

继承 torch.utils.data.Dataset，实现 __len__ 和 __getitem__。

python
from torch.utils.data import Dataset, DataLoader
import torch

class CustomDataset(Dataset):
    def __init__(self, num_samples=1000):
        self.data = torch.randn(num_samples, 10)
        self.labels = torch.randint(0, 2, (num_samples,))

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        return self.data[idx], self.labels[idx]

# 创建数据集
dataset = CustomDataset(2000)
print(f"数据集大小: {len(dataset)}")
print(f"第0个样本: {dataset[0]}")

2.2.2 DataLoader 批量加载 ​

python
dataloader = DataLoader(
    dataset,
    batch_size=32,
    shuffle=True,      # 打乱数据
    num_workers=0,     # 加载进程数（Windows推荐0）
    drop_last=True     # 丢弃最后一个不足 batch_size 的批次
)

for batch_idx, (data, labels) in enumerate(dataloader):
    print(f"批次 {batch_idx}: 数据形状 {data.shape}, 标签形状 {labels.shape}")
    if batch_idx == 2:  # 仅演示前3个批次
        break

2.2.3 实际案例：图像数据集 ​

python
from torchvision import datasets, transforms

# 数据增强与预处理
transform = transforms.Compose([
    transforms.Resize((32, 32)),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

# 使用内置数据集
train_dataset = datasets.CIFAR10(
    root='./data', train=True, download=True, transform=transform
)
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)

2.3 模型定义与训练循环 ​

2.3.1 定义模型（nn.Module） ​

python
import torch.nn as nn
import torch.nn.functional as F

class SimpleNN(nn.Module):
    def __init__(self, input_dim=784, hidden_dim=256, output_dim=10):
        super().__init__()
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_dim, output_dim)

    def forward(self, x):
        x = x.view(x.size(0), -1)  # 展平
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x

model = SimpleNN()
print(model)

2.3.2 完整训练/验证/测试循环 ​

python
import torch.optim as optim

# 超参数
epochs = 5
batch_size = 64
learning_rate = 0.001

# 数据
from torch.utils.data import TensorDataset

X = torch.randn(1000, 784)
y = torch.randint(0, 10, (1000,))

dataset = TensorDataset(X, y)
train_size = int(0.8 * len(dataset))
val_size = len(dataset) - train_size
train_dataset, val_dataset = torch.utils.data.random_split(
    dataset, [train_size, val_size]
)

train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=batch_size)

# 模型、损失、优化器
model = SimpleNN()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)

# 训练循环
for epoch in range(epochs):
    # === 训练阶段 ===
    model.train()
    train_loss = 0.0
    for X_batch, y_batch in train_loader:
        optimizer.zero_grad()
        outputs = model(X_batch)
        loss = criterion(outputs, y_batch)
        loss.backward()
        optimizer.step()
        train_loss += loss.item()

    # === 验证阶段 ===
    model.eval()
    val_loss = 0.0
    correct = 0
    total = 0
    with torch.no_grad():
        for X_batch, y_batch in val_loader:
            outputs = model(X_batch)
            loss = criterion(outputs, y_batch)
            val_loss += loss.item()
            _, predicted = torch.max(outputs, 1)
            total += y_batch.size(0)
            correct += (predicted == y_batch).sum().item()

    print(f"Epoch {epoch+1}/{epochs} | "
          f"Train Loss: {train_loss/len(train_loader):.4f} | "
          f"Val Loss: {val_loss/len(val_loader):.4f} | "
          f"Val Acc: {100*correct/total:.2f}%")

print("训练完成！")

2.4 GPU 加速 ​

2.4.1 设备检测与数据迁移 ​

python
# 检测可用设备
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"使用设备: {device}")

# 将模型和数据移到 GPU
model = SimpleNN().to(device)
x = torch.randn(32, 784).to(device)
y = model(x)
print(y.device)  # cuda:0（如果有 GPU）

2.4.2 多 GPU 并行（DataParallel） ​

python
if torch.cuda.device_count() > 1:
    print(f"使用 {torch.cuda.device_count()} 块 GPU")
    model = nn.DataParallel(model)
model.to(device)

2.4.3 完整 GPU 训练模式 ​

python
def train_gpu(model, train_loader, val_loader, epochs=10, lr=0.001):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = model.to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=lr)

    for epoch in range(epochs):
        model.train()
        for X, y in train_loader:
            X, y = X.to(device), y.to(device)  # 移到 GPU
            optimizer.zero_grad()
            loss = criterion(model(X), y)
            loss.backward()
            optimizer.step()
        print(f"Epoch {epoch+1} 完成")

2.5 模型保存/加载与 ONNX 导出 ​

2.5.1 保存与加载模型 ​

python
# 方式1：保存完整模型（不推荐，依赖于模型定义）
torch.save(model, 'model_full.pth')
model_loaded = torch.load('model_full.pth')

# 方式2：仅保存状态字典（推荐）
torch.save(model.state_dict(), 'model_state_dict.pth')

# 加载状态字典
model = SimpleNN()
model.load_state_dict(torch.load('model_state_dict.pth'))
model.eval()  # 推理模式

2.5.2 检查点（Checkpoint）保存 ​

python
# 保存检查点（包含优化器状态，可从中断处恢复训练）
checkpoint = {
    'epoch': epoch,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'loss': loss,
}
torch.save(checkpoint, 'checkpoint.pth')

# 恢复检查点
checkpoint = torch.load('checkpoint.pth')
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
start_epoch = checkpoint['epoch'] + 1

2.5.3 ONNX 导出（跨平台部署） ​

python
# 导出为 ONNX 格式（可在 ONNX Runtime、TensorRT 等推理）
dummy_input = torch.randn(1, 784).to(device)
torch.onnx.export(
    model,
    dummy_input,
    'model.onnx',
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}}
)

print("ONNX 模型已导出！")

# 使用 ONNX Runtime 推理（需安装 onnxruntime）
# import onnxruntime as ort
# session = ort.InferenceSession('model.onnx')
# result = session.run(None, {'input': dummy_input.cpu().numpy()})

小结 ​

概念关键要点
张量PyTorch 基础数据结构，支持 GPU 与自动微分
Autograd自动构建计算图并计算梯度
Dataset/DataLoader自定义数据加载管道
nn.Module模型定义基类，实现 forward
训练循环train/val/test 三阶段分离
GPU 加速.to(device) 迁移模型和数据
模型保存state_dict 推荐，支持 Checkpoint 断点续训
ONNX跨平台部署格式

练习 ​

  1. 张量操作：创建一个形状为 (3, 224, 224) 的随机张量，将其展平为 (3, -1)，再恢复原形状。
  2. 自定义 Dataset：编写一个读取 CSV 文件的 CSVDataset 类，返回特征和标签。
  3. 完整训练：用 PyTorch 实现一个 3 层 MLP 在 sklearn.datasets.make_moons 上的二分类训练。
  4. GPU 加速：比较同一模型在 CPU 和 GPU（如果有）上的训练时间。
  5. 模型导出：训练好的模型导出为 ONNX 格式，并用 ONNX Runtime 进行推理验证。

Python 学习资料

仅供学习使用