PyTorch 核心 API 完全手册：从基础张量到模型部署

PyTorch 核心 API 完全手册：从基础张量到模型部署

在 PyTorch 的世界里，所有的算法最终都会转化为对张量的各种数学运算。掌握了以下这些 API，你就能看懂 90% 的深度学习源码。

一、 张量创建与基础操作 (Tensor Creation)

1. torch.tensor() - 基础创建

import torch
import numpy as np

# 从 Python 列表创建
t = torch.tensor([1, 2, 3], dtype=torch.float32)

2. torch.from_numpy() - 内存共享转换

arr = np.array([1, 2, 3])
t = torch.from_numpy(arr) # 修改 arr，t 也会随之改变

3. torch.ones() / torch.zeros() - 快速初始化

t_ones = torch.ones(2, 3)  # 全 1 矩阵
t_zeros = torch.zeros(5)   # 全 0 向量

4. torch.randn() - 随机分布

# 生成 3x3 的张量，元素符合标准正态分布 (均值0, 方差1)
t = torch.randn(3, 3)

5. torch.arange() - 序列生成

# 生成 [0, 2, 4, 6, 8]
t = torch.arange(0, 10, 2)

---

二、 维度变换与内存管理 (Shape Manipulation)

6. tensor.view() - 维度重构

t = torch.randn(4, 4)
# 展平为 1D
t_flat = t.view(16)
# 自动推导维度：改为 2 行，列数自动计算
t_auto = t.view(2, -1) 

7. tensor.unsqueeze() / squeeze() - 维度增减

t = torch.randn(3, 3) 
# 在 0 维增加一维，变为 [1, 3, 3]，常用于增加 Batch 维度
t_expanded = t.unsqueeze(0) 

8. torch.cat() - 张量拼接

t1 = torch.randn(2, 3)
t2 = torch.randn(2, 3)
# 在第一维度拼接，变为 [4, 3]
res = torch.cat([t1, t2], dim=0)

9. tensor.permute() - 维度重排

# 图像常见操作：从 [Batch, Height, Width, Channel] 转为 PyTorch 标准 [B, C, H, W]
t = torch.randn(1, 224, 224, 3)
t_p = t.permute(0, 3, 1, 2)

10. tensor.to() - 设备转移

device = "cuda" if torch.cuda.is_available() else "cpu"
t = torch.randn(3).to(device)

---

三、 神经网络层 (NN Layers)

11. nn.Linear() - 全连接层

import torch.nn as nn
fc = nn.Linear(in_features=128, out_features=10)
output = fc(torch.randn(1, 128))

12. nn.Conv2d() - 二维卷积层

# 输入通道3, 输出通道16, 卷积核大小3x3
conv = nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1)
output = conv(torch.randn(1, 3, 32, 32))

13. nn.ReLU() - 激活函数

relu = nn.ReLU()
t = torch.tensor([-1.0, 2.0])
print(relu(t)) # 输出 [0.0, 2.0]

14. nn.MaxPool2d() - 最大池化层

pool = nn.MaxPool2d(kernel_size=2)
# 输入 2x2 区域取最大值，特征图长宽减半

15. nn.Sequential() - 容器

model = nn.Sequential(
    nn.Linear(10, 20),
    nn.ReLU(),
    nn.Linear(20, 2)
)

---

四、 自动求导与优化 (Autograd & Optimization)

16. tensor.requires_grad_() - 追踪梯度

x = torch.randn(3, requires_grad=True)
y = (x * 2).sum()
y.backward() # 反向传播
print(x.grad) # 查看梯度

17. optimizer.zero_grad() - 梯度清零

optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
optimizer.zero_grad() # 每一轮训练开始前必须调用

18. optimizer.step() - 参数更新

# 在 backward() 之后调用，根据梯度更新模型权重
optimizer.step()

19. nn.CrossEntropyLoss() - 损失函数

criterion = nn.CrossEntropyLoss()
logits = torch.randn(3, 5) # 3个样本，5个类别
target = torch.empty(3, dtype=torch.long).random_(5)
loss = criterion(logits, target)

20. torch.no_grad() - 禁用梯度

with torch.no_grad():
    # 测试集预测，不计算梯度以节省内存
    pred = model(test_data)

---

五、 实战：构建一个线性回归模型

import torch
import torch.nn as nn

# 1. 生成数据 (API: randn)
x = torch.randn(100, 1) * 10
y = x + 2 * torch.randn(100, 1)

# 2. 定义模型与优化器 (API: Linear, SGD, MSELoss)
model = nn.Linear(1, 1)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
criterion = nn.MSELoss()

# 3. 训练循环 (API: zero_grad, backward, step)
for epoch in range(100):
    pred = model(x)
    loss = criterion(pred, y)
    
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

# 4. 模型保存 (API: save)
torch.save(model.state_dict(), "model.pth")
print(f"Final Loss: {loss.item()}")

---

六、 进阶：常见报错排查与排除

21. RuntimeError: Expected object of device type cuda but got cpu

• 原因：模型在 GPU 上，但数据还在 CPU 上（或反之）。
• 解决 ：确保 model.to(device) 和 data.to(device) 指向同一个设备。

22. RuntimeError: size mismatch

• 原因 ：nn.Linear 的 in_features 与上一层的输出维度不匹配。
• 调试 ：在 forward 中加入 print(x.shape) 查看进入 Linear 层前的维度。

23. tensor.detach() 与 tensor.cpu().numpy()

• 技巧 ：在绘图或输出结果时，如果张量在 GPU 上且有梯度，必须先 detach() 再转为 NumPy。

# 正确写法
plt.plot(t.detach().cpu().numpy())

---

相关知识点讲解：什么是动态计算图？

PyTorch 的核心是 DAG (Directed Acyclic Graph)。

• 有向 (Directed)：计算有明确的方向，从输入到输出。
• 无向 (Acyclic)：图中没有回路。
• 动态：每当你执行一次代码，PyTorch 就会重新构建一张图。这允许你在训练过程中根据条件判断改变网络的结构（例如，如果损失值大于某个阈值，就多走一个全连接层）。

---

AI创作声明：本文部分内容由 AI 辅助生成，并经人工整理与验证，仅供参考学习，欢迎指出错误与不足之处。

---