PyTorch 深度学习框架核心函数与实战指南

3. 索引与切片

t = torch.arange(10).reshape(2, 5)
print(t[0, :])      # 第一行
print(t[:, 1])      # 第二列
print(t[1:3, :])    # 切片

三、自动求导系统（Autograd）

PyTorch 的动态计算图使得反向传播变得非常直观。

1. 开启梯度追踪

x = torch.tensor([1.0, 2.0], requires_grad=True)
y = x ** 2
y.sum().backward()
print(x.grad) # 输出梯度

2. 停止梯度计算

with torch.no_grad():
    z = x * 2

四、构建神经网络模型

PyTorch 通过 nn.Module 类来定义网络结构。

1. 定义网络层

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

class SimpleNet(nn.Module):
    def __init__(self):
        super(SimpleNet, self).__init__()
        self.fc1 = nn.Linear(784, 128)
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(p=0.5)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = x.view(-1, 784)
        x = self.fc1(x)
        x = self.relu(x)
        x = self.dropout(x)
        x = self.fc2(x)
        return x

2. 激活函数详解

• Sigmoid: 将输出压缩到 (0, 1)，常用于二分类。
• Tanh: 将输出压缩到 (-1, 1)，解决 Sigmoid 梯度消失问题。
• ReLU: Rectified Linear Unit，目前最常用的激活函数，计算高效。
• Softmax: 用于多分类输出的概率分布。

3. 损失函数

criterion = nn.CrossEntropyLoss() # 分类任务
# criterion = nn.MSELoss()       # 回归任务

五、数据加载与预处理

1. Dataset 与 DataLoader

from torch.utils.data import Dataset, DataLoader

class MyDataset(Dataset):
    def __init__(self, data, labels):
        self.data = data
        self.labels = labels

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

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

dataset = MyDataset(data, labels)
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

六、训练循环（Training Loop）

完整的训练流程包括前向传播、计算损失、反向传播和优化器更新。

model = SimpleNet()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

for epoch in range(10):
    for inputs, labels in dataloader:
        # 前向传播
        outputs = model(inputs)
        loss = criterion(outputs, labels)

        # 反向传播
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    print(f"Epoch {epoch}, Loss: {loss.item()}")

七、优化器详解

• SGD (Stochastic Gradient Descent): 基础优化算法，需配合动量使用。
• Adam: 自适应学习率，通常效果优于 SGD，收敛更快。
• RMSprop: 适用于 RNN 等序列模型。

optimizer_sgd = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
optimizer_adam = torch.optim.Adam(model.parameters(), lr=0.001)

八、模型保存与加载

# 保存整个模型
torch.save(model.state_dict(), 'checkpoint.pth')

# 加载模型
model.load_state_dict(torch.load('checkpoint.pth'))

九、主流框架对比

十、实战建议与总结

1. 从零开始：理解底层数学原理，尝试用 NumPy 实现简单的神经网络。
2. 框架选择：科研推荐 PyTorch，生产部署考虑 TensorFlow Lite 或 ONNX。
3. 代码规范：遵循 PEP8，模块化设计，便于维护。
4. 资源利用：善用 GPU 加速，注意显存管理，避免 OOM。

通过掌握上述内容，开发者可以独立搭建和设计卷积神经网络，并进行神经网络的训练和推理。无论是分类、检测还是生成任务，PyTorch 都能提供强大的支持。建议结合官方文档和开源项目持续深入学习，提升工程实践能力。