PyTorch API 详细中文文档 按模块分类，涵盖核心函数与用法示例
目录 张量操作 (Tensor Operations) 数学运算 (Math Operations) 自动求导 (Autograd) 神经网络模块 (torch.nn) 优化器 (torch.optim) 数据加载与处理 (torch.utils.data) 设备管理 (Device Management) 模型保存与加载 分布式训练 (Distributed Training) 实用工具函数 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 + ytorch.mul(x, y)乘法 torch.mul(x, y) 或 x * ytorch.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.grad → 4.0（若 y = x²）torch.no_grad()禁用梯度追踪 with torch.no_grad(): y = x * 2detach()分离张量（不追踪梯度） y = x.detach()
3.2 梯度控制 函数 描述 x.retain_grad()保留非叶子节点的梯度 torch.autograd.grad(outputs, inputs)手动计算梯度
示例 ：
x = torch. tensor( 3.0 , requires_grad= True )
y = x** 3 + 2 * x
dy_dx = torch. autograd. grad( y, x)
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)
自定义数据集示例 ：
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) 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)
常见问题与技巧 GPU 内存不足
使用 batch_size 较小的值 启用混合精度训练 (torch.cuda.amp) 使用 torch.utils.checkpoint 节省内存 梯度爆炸/消失
使用梯度裁剪：torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) 调整权重初始化方法 模型推理模式
model. eval ( )
with torch. no_grad( ) : outputs = model( inputs)
文档说明
本文档基于 PyTorch 2.5 编写，部分 API 可能不兼容旧版本。 更详细的参数说明请参考 PyTorch 官方文档。 
