PyTorch 作为当前主流的深度学习框架之一，凭借其动态图机制和灵活的 API 设计，在学术界与工业界均占据重要地位。本文系统梳理了 PyTorch 的核心功能模块，涵盖张量操作、自动求导、神经网络构建及训练流程等关键知识点。通过对比 TensorFlow 与 Keras，分析 PyTorch 的优势，并提供从零搭建卷积神经网络的实战示例。内容涉及激活函数详解、优化器选择策略及多框架协同应用，旨在帮助开发者快速掌握深度学习开源框架，提升模型训练与推理能力。
PyTorch 是目前常用的深度学习框架之一，它凭借着对初学者的友好性、灵活性，发展迅猛。相比于 TensorFlow 的静态图限制和 Keras 的高度封装，PyTorch 无论是在学术圈还是工业界，都相当占优势。掌握了 PyTorch，就相当于走上了深度学习、机器学习的快车道。
本文旨在系统梳理 PyTorch 的核心功能模块，涵盖张量操作、自动求导、神经网络构建及训练流程等关键知识点，帮助开发者快速掌握该框架。
在使用 PyTorch 之前，需要确保 Python 环境已正确配置。推荐使用 Anaconda 管理虚拟环境。
conda create -n pytorch_env python=3.8
conda activate pytorch_env
根据操作系统和 CUDA 版本选择对应的安装命令。例如，使用 NVIDIA GPU 加速：
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
若无 GPU，可安装 CPU 版本：
pip install torch torchvision torchaudio
张量是 PyTorch 中的核心数据结构，类似于 NumPy 的 ndarray，但支持在 GPU 上运行。
import torch
# 从列表创建
tensor_list = torch.tensor([[1, 2], [3, 4]])
# 随机初始化
random_tensor = torch.rand(2, 3)
# 全零/全一
zeros_tensor = torch.zeros(2, 3)
one_tensor = torch.ones(2, 3)
# 指定数据类型和位置
tensor_cpu = torch.tensor([1.0, 2.0])
tensor_gpu = tensor_cpu.cuda() if torch.cuda.is_available() else tensor_cpu
a = torch.tensor([1.0, 2.0, 3.0])
b = torch.tensor([4.0, 5.0, 6.0])
# 加减乘除
c_add = a + b
c_mul = a * b
c_div = a / b
# 矩阵乘法
dot_product = torch.matmul(a, b) # 若为向量则返回标量
matrix_a = torch.rand(2, 3)
matrix_b = torch.rand(3, 2)
matmul_result = torch.mm(matrix_a, matrix_b)
t = torch.arange(10).reshape(2, 5)
print(t[0, :]) # 第一行
print(t[:, 1]) # 第二列
print(t[1:3, :]) # 切片
PyTorch 的动态计算图使得反向传播变得非常直观。
x = torch.tensor([1.0, 2.0], requires_grad=True)
y = x ** 2
y.sum().backward()
print(x.grad) # 输出梯度
with torch.no_grad():
z = x * 2
PyTorch 通过 nn.Module 类来定义网络结构。
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
criterion = nn.CrossEntropyLoss() # 分类任务
# criterion = nn.MSELoss() # 回归任务
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)
完整的训练流程包括前向传播、计算损失、反向传播和优化器更新。
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()}")
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'))
| 特性 | PyTorch | TensorFlow | Keras |
|---|---|---|---|
| 图机制 | 动态图 | 静态图 (TF1) / 动态 (TF2) | 基于 TF 的高层封装 |
| 调试 | 易于调试 (Pythonic) | 较复杂 | 简单 |
| 社区 | 学术界主导 | 工业界主导 | 易用性高 |
| 部署 | TorchScript | TFLite | ONNX |
通过掌握上述内容，开发者可以独立搭建和设计卷积神经网络，并进行神经网络的训练和推理。无论是分类、检测还是生成任务，PyTorch 都能提供强大的支持。建议结合官方文档和开源项目持续深入学习，提升工程实践能力。
微信公众号「极客日志」，在微信中扫描左侧二维码关注。展示文案：极客日志 zeeklog
使用加密算法（如AES、TripleDES、Rabbit或RC4）加密和解密文本明文。 在线工具，加密/解密文本在线工具，online
生成新的随机RSA私钥和公钥pem证书。 在线工具，RSA密钥对生成器在线工具，online
基于 Mermaid.js 实时预览流程图、时序图等图表，支持源码编辑与即时渲染。 在线工具，Mermaid 预览与可视化编辑在线工具，online
解析常见 curl 参数并生成 fetch、axios、PHP curl 或 Python requests 示例代码。 在线工具，curl 转代码在线工具，online
将字符串编码和解码为其 Base64 格式表示形式即可。 在线工具，Base64 字符串编码/解码在线工具，online
将字符串、文件或图像转换为其 Base64 表示形式。 在线工具，Base64 文件转换器在线工具，online