PyTorch API 详细中文文档,基于PyTorch2.5


PyTorch API 详细中文文档

按模块分类,涵盖核心函数与用法示例


目录

  1. 张量操作 (Tensor Operations)
  2. 数学运算 (Math Operations)
  3. 自动求导 (Autograd)
  4. 神经网络模块 (torch.nn)
  5. 优化器 (torch.optim)
  6. 数据加载与处理 (torch.utils.data)
  7. 设备管理 (Device Management)
  8. 模型保存与加载
  9. 分布式训练 (Distributed Training)
  10. 实用工具函数

1. 张量操作 (Tensor Operations)

1.1 张量创建
函数 描述 示例
torch.tensor(data, dtype, device) 从数据创建张量 torch.tensor([1,2,3], dtype=torch.float32)
torch.zeros(shape) 创建全零张量 torch.zeros(2,3)
torch.ones(shape) 创建全一张量 torch.ones(5)
torch.rand(shape) 均匀分布随机张量 torch.rand(3,3)
torch.randn(shape) 标准正态分布张量 torch.randn(4,4)
torch.arange(start, end, step) 创建等差序列 torch.arange(0, 10, 2)[0,2,4,6,8]
torch.linspace(start, end, steps) 线性间隔序列 torch.linspace(0, 1, 5)[0, 0.25, 0.5, 0.75, 1]
1.2 张量属性
属性/方法 描述 示例
.shape 张量维度 x = torch.rand(2,3); x.shape → torch.Size([2,3])
.dtype 数据类型 x.dtype → torch.float32
.device 所在设备 x.device → device(type='cpu')
.requires_grad 是否追踪梯度 x.requires_grad = True
1.3 张量变形
函数 描述 示例
.view(shape) 调整形状(不复制数据) x = torch.arange(6); x.view(2,3)
.reshape(shape) 类似 view,但自动处理内存连续性 x.reshape(3,2)
.permute(dims) 调整维度顺序 x = torch.rand(2,3,4); x.permute(1,2,0)
.squeeze(dim) 去除大小为1的维度 x = torch.rand(1,3); x.squeeze(0)shape [3]
.unsqueeze(dim) 添加大小为1的维度 x = torch.rand(3); x.unsqueeze(0)shape [1,3]

2. 数学运算 (Math Operations)

2.1 逐元素运算
函数 描述 示例
torch.add(x, y) 加法 torch.add(x, y)x + y
torch.mul(x, y) 乘法 torch.mul(x, y)x * y
torch.exp(x) 指数运算 torch.exp(torch.tensor([1.0]))[2.7183]
torch.log(x) 自然对数 torch.log(torch.exp(tensor([2.0])))[2.0]
torch.clamp(x, min, max) 限制值范围 torch.clamp(x, min=0, max=1)
2.2 矩阵运算
函数 描述 示例
torch.matmul(x, y) 矩阵乘法 x = torch.rand(2,3); y = torch.rand(3,4); torch.matmul(x, y)
torch.inverse(x) 矩阵求逆 x = torch.rand(3,3); inv_x = torch.inverse(x)
torch.eig(x) 特征值分解 eigenvalues, eigenvectors = torch.eig(x)
2.3 统计运算
函数 描述 示例
torch.sum(x, dim) 沿维度求和 x = torch.rand(2,3); torch.sum(x, dim=1)
torch.mean(x, dim) 沿维度求均值 torch.mean(x, dim=0)
torch.max(x, dim) 沿维度求最大值 values, indices = torch.max(x, dim=1)
torch.argmax(x, dim) 最大值索引 indices = torch.argmax(x, dim=1)

3. 自动求导 (Autograd)

3.1 梯度计算
函数/属性 描述 示例
x.backward() 反向传播计算梯度 x = torch.tensor(2.0, requires_grad=True); y = x**2; y.backward()
x.grad 查看梯度值 x.grad4.0(若 y = x²
torch.no_grad() 禁用梯度追踪 with torch.no_grad(): y = x * 2
detach() 分离张量(不追踪梯度) y = x.detach()
3.2 梯度控制
函数 描述
x.retain_grad() 保留非叶子节点的梯度
torch.autograd.grad(outputs, inputs) 手动计算梯度

示例

python 复制代码
x = torch.tensor(3.0, requires_grad=True)  
y = x**3 + 2*x  
dy_dx = torch.autograd.grad(y, x)  # 返回 (torch.tensor(29.0),)  

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

4.1 层定义
描述 示例
nn.Linear(in_features, out_features) 全连接层 layer = nn.Linear(784, 256)
nn.Conv2d(in_channels, out_channels, kernel_size) 卷积层 conv = nn.Conv2d(3, 16, kernel_size=3)
nn.LSTM(input_size, hidden_size) LSTM 层 lstm = nn.LSTM(100, 50)
nn.Dropout(p=0.5) Dropout 层 dropout = nn.Dropout(0.2)
4.2 激活函数
函数 描述 示例
nn.ReLU() ReLU 激活 F.relu(x)nn.ReLU()(x)
nn.Sigmoid() Sigmoid 函数 torch.sigmoid(x)
nn.Softmax(dim) Softmax 归一化 F.softmax(x, dim=1)
4.3 损失函数
描述 示例
nn.MSELoss() 均方误差 loss_fn = nn.MSELoss()
nn.CrossEntropyLoss() 交叉熵损失 loss = loss_fn(outputs, labels)
nn.BCELoss() 二分类交叉熵 loss_fn = nn.BCELoss()

5. 优化器 (torch.optim)

5.1 优化器定义
描述 示例
optim.SGD(params, lr) 随机梯度下降 optimizer = optim.SGD(model.parameters(), lr=0.01)
optim.Adam(params, lr) Adam 优化器 optimizer = optim.Adam(model.parameters(), lr=0.001)
optim.RMSprop(params, lr) RMSprop 优化器 optimizer = optim.RMSprop(params, lr=0.01)
5.2 优化器方法
方法 描述 示例
optimizer.zero_grad() 清空梯度 optimizer.zero_grad()
optimizer.step() 更新参数 loss.backward(); optimizer.step()
optimizer.state_dict() 获取优化器状态 state = optimizer.state_dict()

6. 数据加载与处理 (torch.utils.data)

6.1 数据集类
类/函数 描述 示例
Dataset 自定义数据集基类 继承并实现 __len____getitem__
DataLoader(dataset, batch_size, shuffle) 数据加载器 loader = DataLoader(dataset, batch_size=64, shuffle=True)

自定义数据集示例

python 复制代码
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]  
6.2 数据预处理 (TorchVision)
python 复制代码
from torchvision import transforms  

transform = transforms.Compose([  
    transforms.Resize(256),          # 调整图像大小  
    transforms.ToTensor(),           # 转为张量  
    transforms.Normalize(mean=[0.5], std=[0.5])  # 标准化  
])  

7. 设备管理 (Device Management)

7.1 设备切换
函数/方法 描述 示例
.to(device) 移动张量/模型到设备 x = x.to('cuda:0')
torch.cuda.is_available() 检查 GPU 是否可用 if torch.cuda.is_available(): ...
torch.cuda.empty_cache() 清空 GPU 缓存 torch.cuda.empty_cache()

8. 模型保存与加载

函数 描述 示例
torch.save(obj, path) 保存对象(模型/参数) torch.save(model.state_dict(), 'model.pth')
torch.load(path) 加载对象 model.load_state_dict(torch.load('model.pth'))
model.state_dict() 获取模型参数字典 params = model.state_dict()

9. 分布式训练 (Distributed Training)

函数/类 描述 示例
nn.DataParallel(model) 单机多卡并行 model = nn.DataParallel(model)
torch.distributed.init_process_group() 初始化分布式训练 需配合多进程使用

10. 实用工具函数

函数 描述 示例
torch.cat(tensors, dim) 沿维度拼接张量 torch.cat([x, y], dim=0)
torch.stack(tensors, dim) 堆叠张量(新建维度) torch.stack([x, y], dim=1)
torch.split(tensor, split_size, dim) 分割张量 chunks = torch.split(x, 2, dim=0)

常见问题与技巧

  1. GPU 内存不足

    • 使用 batch_size 较小的值
    • 启用混合精度训练 (torch.cuda.amp)
    • 使用 torch.utils.checkpoint 节省内存
  2. 梯度爆炸/消失

    • 使用梯度裁剪:torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
    • 调整权重初始化方法
  3. 模型推理模式

    python 复制代码
    model.eval()  # 关闭 Dropout 和 BatchNorm 的随机性  
    with torch.no_grad():  
        outputs = model(inputs)  

文档说明

  • 本文档基于 PyTorch 2.5 编写,部分 API 可能不兼容旧版本。
  • 更详细的参数说明请参考 PyTorch 官方文档
相关推荐
Monkey-旭4 分钟前
Android Bitmap 完全指南:从基础到高级优化
android·java·人工智能·计算机视觉·kotlin·位图·bitmap
天若有情6731 小时前
【python】Python爬虫入门教程:使用requests库
开发语言·爬虫·python·网络爬虫·request
哪 吒2 小时前
OpenAI放大招:ChatGPT学习模式上线,免费AI智能家教
人工智能·学习·ai·chatgpt·gemini·deepseek
IT北辰2 小时前
用Python+MySQL实战解锁企业财务数据分析
python·mysql·数据分析
Lucky高2 小时前
selenium(WEB自动化工具)
python
老鱼说AI2 小时前
循环神经网络RNN原理精讲,详细举例!
人工智能·rnn·深度学习·神经网络·自然语言处理·语音识别
秃然想通2 小时前
掌握Python三大语句:顺序、条件与循环
开发语言·python·numpy
lingling0092 小时前
颐顿机电携手观远BI数据:以数据驱动决策,领跑先进制造智能化升级
大数据·人工智能·制造
b***25112 小时前
电池自动生产线:科技赋能下的高效制造新范式
大数据·人工智能
EVERSPIN2 小时前
分享低功耗单火线开关语音识别方案
人工智能·语音识别