PyTorch 常用方法与函数

PyTorch 是目前最流行的深度学习框架之一。本文档总结了 torch 库中使用频率最高的方法和函数，并附带了代码示例，方便快速查阅。

1. 张量创建 (Tensor Creation)

张量（Tensor）是 PyTorch 的核心数据结构，类似于 NumPy 的 ndarray，但支持 GPU 加速。

常用函数

函数	描述
torch.tensor(data)	直接从数据（如列表、NumPy 数组）创建张量
torch.zeros(shape)	创建全 0 张量
torch.ones(shape)	创建全 1 张量
torch.randn(shape)	创建符合标准正态分布（均值0，方差1）的随机张量
torch.rand(shape)	创建 [0, 1) 区间均匀分布的随机张量
torch.arange(start, end, step)	创建等差数列张量
torch.from_numpy(ndarray)	将 NumPy 数组转换为 Tensor（共享内存）

代码示例

import torch
import numpy as np

# 1. 直接创建
x = torch.tensor([[1, 2], [3, 4]])
print(f"Tensor:\n{x}")

# 2. 随机初始化 (3行4列)
rand_x = torch.randn(3, 4)

# 3. 全0和全1
zeros = torch.zeros(2, 2, dtype=torch.long) # 指定数据类型
ones = torch.ones(2, 2)

# 4. Numpy 互转
np_arr = np.array([1, 2, 3])
tensor_from_np = torch.from_numpy(np_arr)
back_to_np = tensor_from_np.numpy()

2. 张量属性与维度操作 (Shape & Manipulation)

深度学习中，90% 的 Bug 通常来自于维度（Shape）不匹配。

常用函数

函数	描述
tensor.shape / tensor.size()	获取张量形状
tensor.view(shape)	改变形状（需数据连续），相当于 Reshape
tensor.reshape(shape)	改变形状（更通用，必要时会复制数据）
tensor.unsqueeze(dim)	在指定位置增加一个维度
tensor.squeeze(dim)	压缩指定位置为 1 的维度
tensor.permute(dims)	交换维度顺序（类似于转置，但支持多维）
torch.cat(tensors, dim)	在指定维度拼接多个张量
torch.stack(tensors, dim)	在新维度堆叠多个张量

代码示例

x = torch.randn(4, 3) # Shape: [4, 3]

# 1. 改变形状 [4, 3] -> [12]
flat_x = x.view(-1) 

# 2. 增加维度 [4, 3] -> [4, 1, 3]
expanded_x = x.unsqueeze(1)

# 3. 维度交换 [4, 3] -> [3, 4]
permuted_x = x.permute(1, 0)

# 4. 拼接
t1 = torch.randn(2, 3)
t2 = torch.randn(2, 3)
# 按行拼接 -> [4, 3]
cat_res = torch.cat([t1, t2], dim=0) 

3. 数学运算 (Math Operations)

PyTorch 支持丰富的数学运算，支持广播机制（Broadcasting）。

常用函数

函数	描述
torch.add(x, y) / x + y	加法
torch.mul(x, y) / x * y	元素级乘法 (Element-wise)
torch.matmul(x, y) / x @ y	矩阵乘法 (Matrix Multiplication)
torch.sum(x, dim)	求和
torch.mean(x, dim)	求平均值
torch.max(x, dim)	求最大值（返回 values 和 indices）
torch.argmax(x, dim)	返回最大值的索引

代码示例

x = torch.tensor([[1., 2.], [3., 4.]])
y = torch.tensor([[2., 2.], [2., 2.]])

# 1. 元素乘法
print(x * y) 
# 输出: [[2., 4.], [6., 8.]]

# 2. 矩阵乘法
print(x @ y) 
# 输出: [[6., 6.], [14., 14.]]

# 3. 聚合操作
print(x.sum())       # 所有元素求和
print(x.mean(dim=0)) # 按列求均值 -> [2., 3.]

4. 自动求导 (AutoGrad)

PyTorch 的核心功能，用于反向传播计算梯度。

核心概念

• requires_grad=True: 告诉 PyTorch 追踪该张量的所有操作。
• backward(): 自动计算梯度。
• x.grad: 存储计算出的梯度值。
• torch.no_grad(): 上下文管理器，用于推理阶段，停止梯度计算以节省显存。

代码示例

# 创建需计算梯度的张量
w = torch.tensor([2.0], requires_grad=True)
x = torch.tensor([3.0])
b = torch.tensor([1.0], requires_grad=True)

# 前向传播: y = w * x + b
y = w * x + b  # y = 2*3 + 1 = 7

# 反向传播
y.backward()

# 查看梯度
# dy/dw = x = 3
print(f"w.grad: {w.grad}") 
# dy/db = 1
print(f"b.grad: {b.grad}") 

5. 神经网络模块 (torch.nn)

torch.nn 提供了构建神经网络所需的层。

常用层

类	描述
nn.Linear(in, out)	全连接层 (Dense Layer)
nn.Conv2d(in, out, k)	2D 卷积层 (用于图像)
nn.LSTM(in, hidden)	循环神经网络层
nn.ReLU() / nn.Sigmoid()	激活函数
nn.Dropout(p)	随机失活（防止过拟合）
nn.BatchNorm2d(features)	批归一化

常用损失函数

• nn.MSELoss(): 均方误差（用于回归）。
• nn.CrossEntropyLoss(): 交叉熵损失（用于分类）。

代码示例

import torch.nn as nn

class SimpleNet(nn.Module):
    def __init__(self):
        super(SimpleNet, self).__init__()
        # 定义层
        self.fc1 = nn.Linear(10, 50)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(50, 2)
    
    def forward(self, x):
        # 定义前向传播路径
        out = self.fc1(x)
        out = self.relu(out)
        out = self.fc2(out)
        return out

model = SimpleNet()
input_data = torch.randn(5, 10) # Batch size=5, features=10
output = model(input_data)
print(output.shape) # [5, 2]

6. 优化器 (torch.optim)

用于更新模型参数。

常用优化器

• torch.optim.SGD(params, lr, momentum)
• torch.optim.Adam(params, lr)

典型训练步骤 (Boilerplate)

import torch.optim as optim

# 1. 定义模型和优化器
model = SimpleNet()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 2. 模拟训练循环
for epoch in range(10):
    # 假设 data 和 targets 是数据
    data = torch.randn(5, 10)
    targets = torch.empty(5, dtype=torch.long).random_(2)
    
    # [关键步骤 1] 清空梯度
    optimizer.zero_grad()
    
    # [关键步骤 2] 前向传播
    outputs = model(data)
    loss = criterion(outputs, targets)
    
    # [关键步骤 3] 反向传播计算梯度
    loss.backward()
    
    # [关键步骤 4] 更新参数
    optimizer.step()
    
    print(f"Epoch {epoch}, Loss: {loss.item():.4f}")

7. 设备管理 (CPU/GPU)

将张量或模型移动到显卡进行加速。

常用命令

命令	描述
torch.cuda.is_available()	检查 CUDA 是否可用
tensor.to(device)	移动张量到指定设备
model.to(device)	移动模型参数到指定设备
item.cpu()	移回 CPU（如需转 numpy，必须先移回 CPU）

代码示例

# 检测设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 移动张量
x = torch.randn(3, 3).to(device)

# 移动模型
model = SimpleNet().to(device)

print(f"Running on: {device}")

8. 模型保存与加载

函数	描述
torch.save(obj, path)	保存模型或张量
torch.load(path)	加载模型或张量
model.state_dict()	获取模型参数字典（推荐保存方式）
model.load_state_dict(dict)	加载参数字典

代码示例

# 推荐：只保存参数（权重）
torch.save(model.state_dict(), "model_weights.pth")

# 加载参数
model = SimpleNet() # 先实例化结构
model.load_state_dict(torch.load("model_weights.pth"))
model.eval() # 切换到评估模式