PyTorch实战指南：从张量操作到模型训练

01-24 4105阅读

引言：为什么选择PyTorch

PyTorch作为当前最受欢迎的深度学习框架之一，凭借其动态计算图和直观的API设计，在研究和工业界都获得了广泛认可。本实战指南将带你体验PyTorch的核心功能，通过实际代码演示完成从基础操作到完整模型训练的全流程。

PyTorch核心组件实战

张量操作基础

张量（Tensor）是PyTorch中最基本的数据结构，类似于NumPy的多维数组但支持GPU加速：

import torch

# 创建张量
x = torch.tensor([[1, 2], [3, 4]])
y = torch.ones(2, 2)  # 全1矩阵

# 数学运算
z = x + y
print(f"加法结果:\n{z}")

# 设备转移
if torch.cuda.is_available():
    device = torch.device("cuda")
    x_gpu = x.to(device)  # 转移到GPU
    print(f"GPU上的张量:\n{x_gpu}")

自动微分机制

PyTorch的自动求导（Autograd）系统是构建神经网络的核心：

PyTorch实战指南：从张量操作到模型训练

# 创建需要梯度的张量
a = torch.tensor(3.0, requires_grad=True)
b = torch.tensor(2.0, requires_grad=True)

# 定义计算图
c = a**2 + b**3
d = c * 4

# 反向传播计算梯度
d.backward()

# 打印梯度值
print(f"a的梯度: {a.grad}")  # 输出: 24.0 (2*3*4)
print(f"b的梯度: {b.grad}")  # 输出: 24.0 (3*2^2*4)

构建神经网络模型

使用PyTorch的nn.Module类创建简单全连接网络：

import torch.nn as nn

class SimpleNN(nn.Module):
    def __init__(self):
        super(SimpleNN, self).__init__()
        # 定义网络层
        self.fc1 = nn.Linear(784, 128)  # 输入784维，输出128维
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(128, 10)   # 输出10分类

    def forward(self, x):
        # 前向传播路径
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x

# 实例化模型
model = SimpleNN()
print(f"模型结构:\n{model}")

完整训练流程实战

在MNIST数据集上训练分类模型：

from torchvision import datasets, transforms
from torch.utils.data import DataLoader

# 数据预处理
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

# 加载数据集
train_set = datasets.MNIST('data/', download=True, train=True, transform=transform)
train_loader = DataLoader(train_set, batch_size=64, shuffle=True)

# 初始化组件
model = SimpleNN()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
criterion = nn.CrossEntropyLoss()

# 训练循环
for epoch in range(5):  # 5个epoch
    for images, labels in train_loader:
        # 扁平化图像 [64, 1, 28, 28] -> [64, 784]
        images = images.view(images.shape[0], -1)

        # 前向传播
        outputs = model(images)
        loss = criterion(outputs, labels)

        # 反向传播
        optimizer.zero_grad()  # 清除历史梯度
        loss.backward()        # 计算梯度
        optimizer.step()       # 更新参数

    print(f'Epoch [{epoch+1}/5], Loss: {loss.item():.4f}')

模型保存与加载

训练完成后保存模型以备后续使用：

# 保存模型
torch.save(model.state_dict(), 'mnist_model.pth')

# 加载模型
loaded_model = SimpleNN()
loaded_model.load_state_dict(torch.load('mnist_model.pth'))
loaded_model.eval()  # 设置为评估模式