Pytorch torch 参考手册

作者likuolei

PyTorch 是一个开源的机器学习框架，其核心模块 torch 提供了大量用于张量操作、神经网络构建、优化和自动求导的工具。以下是关于 PyTorch torch 模块的参考手册概述，涵盖核心功能、使用方法和关键资源，力求简洁且全面。

1. torch 模块概述

torch 是 PyTorch 的核心库，包含以下主要功能：

  • 张量操作：类似 NumPy 的多维数组操作，支持 GPU 加速。
  • 自动求导：通过 torch.autograd 实现动态计算图和梯度计算。
  • 神经网络模块：torch.nn 提供构建神经网络的工具。
  • 优化器：torch.optim 提供优化算法（如 SGD、Adam）。
  • 数据处理：与 torch.utils.data 配合进行数据加载和预处理。
  • 设备管理：支持 CPU/GPU 切换和分布式计算。

2. 核心组件与使用方法

(1) 张量（Tensor）

torch.Tensor 是 PyTorch 的核心数据结构，类似 NumPy 的 ndarray，支持 GPU 加速。

  • 创建张量
  import torch

  # 基本创建
  x = torch.tensor([1, 2, 3])  # 1D 张量
  y = torch.randn(2, 3)  # 2x3 随机张量
  z = torch.zeros(2, 3)  # 2x3 全零张量
  ones = torch.ones(2, 3)  # 2x3 全一张量
  • 张量操作
  # 基本运算
  add = x + y
  mul = torch.matmul(y, y.t())  # 矩阵乘法

  # 形状操作
  reshaped = y.view(3, 2)  # 重塑形状
  transposed = y.t()  # 转置

  # 设备转换
  if torch.cuda.is_available():
      y = y.to('cuda')  # 移动到 GPU
  • 常用属性
  print(y.shape)  # 张量形状: torch.Size([2, 3])
  print(y.dtype)  # 数据类型: torch.float32
  print(y.device)  # 设备: cpu 或 cuda

(2) 自动求导（Autograd）

torch.autograd 提供自动梯度计算，支持动态计算图。

  • 启用梯度跟踪
  x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
  y = x * 2
  z = y.mean()
  z.backward()  # 反向传播
  print(x.grad)  # 梯度: tensor([0.6667, 0.6667, 0.6667])
  • 禁用梯度
  with torch.no_grad():
      y = x * 2  # 不跟踪梯度，节省内存

(3) 神经网络（torch.nn）

torch.nn 提供构建神经网络的模块化工具。

  • 定义模型
  import torch.nn as nn

  class SimpleNet(nn.Module):
      def __init__(self):
          super(SimpleNet, self).__init__()
          self.fc1 = nn.Linear(10, 5)
          self.fc2 = nn.Linear(5, 2)

      def forward(self, x):
          x = torch.relu(self.fc1(x))
          x = self.fc2(x)
          return x

  model = SimpleNet()
  • 常用层
  • nn.Linear(in_features, out_features)：全连接层
  • nn.Conv2d(in_channels, out_channels, kernel_size)：卷积层
  • nn.ReLU()：激活函数
  • nn.Dropout(p=0.5)：随机丢弃，防止过拟合

(4) 损失函数

torch.nn 提供多种损失函数：

criterion = nn.CrossEntropyLoss()  # 交叉熵损失
mse_loss = nn.MSELoss()  # 均方误差

(5) 优化器（torch.optim）

torch.optim 提供优化算法：

import torch.optim as optim

optimizer = optim.Adam(model.parameters(), lr=0.001)
optimizer.zero_grad()  # 清零梯度
loss = criterion(output, target)
loss.backward()  # 计算梯度
optimizer.step()  # 更新参数
  • 常用优化器
  • optim.SGD：随机梯度下降
  • optim.Adam：自适应矩估计
  • optim.RMSprop：RMSprop 优化器

(6) 数据加载（torch.utils.data）

与数据转换相关，参考前文关于 Dataset 和 DataLoader 的内容。

  • 示例
  from torch.utils.data import DataLoader, TensorDataset

  data = torch.randn(100, 10)
  labels = torch.randint(0, 2, (100,))
  dataset = TensorDataset(data, labels)
  dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

(7) 设备管理

支持 CPU 和 GPU 切换：

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = model.to(device)
data = data.to(device)

3. 常用功能与方法

张量操作

  • 创建：torch.tensor(), torch.randn(), torch.arange(), torch.linspace()
  • 数学运算：torch.add(), torch.mul(), torch.matmul(), torch.sum(), torch.mean()
  • 索引与切片：类似 NumPy，支持 tensor[1:3], tensor[:, 1]
  • 形状操作：tensor.view(), tensor.reshape(), tensor.permute()
  • 类型转换：tensor.float(), tensor.int(), tensor.to(dtype=torch.float32)

自动求导

  • tensor.requires_grad：启用梯度跟踪
  • tensor.backward()：计算梯度
  • torch.no_grad()：禁用梯度计算
  • torch.autograd.grad()：显式计算梯度

神经网络

  • 模块：nn.Module, nn.Sequential
  • 激活函数：nn.ReLU, nn.Sigmoid, nn.Tanh
  • 池化与规范化：nn.MaxPool2d, nn.BatchNorm2d

工具函数

  • torch.save(model.state_dict(), 'model.pth')：保存模型
  • model.load_state_dict(torch.load('model.pth'))：加载模型
  • torch.cuda.is_available()：检查 GPU 可用性

4. 进阶功能

  • 分布式训练：torch.distributed 支持多 GPU 和多节点训练。
  from torch.nn.parallel import DistributedDataParallel as DDP
  model = DDP(model)
  • 自定义运算：通过 torch.autograd.Function 定义自定义前向/反向传播。
  • ONNX 导出：将模型导出为 ONNX 格式以跨框架使用。
  torch.onnx.export(model, dummy_input, 'model.onnx')
  • 混合精度训练：使用 torch.cuda.amp 提高性能。
  from torch.cuda.amp import autocast, GradScaler
  scaler = GradScaler()
  with autocast():
      output = model(input)
      loss = criterion(output, target)
  scaler.scale(loss).backward()
  scaler.step(optimizer)
  scaler.update()

5. 常见问题与注意事项

  • 张量设备一致性：确保模型和数据在同一设备（CPU/GPU）上。
  • 梯度累积：调用 optimizer.zero_grad() 避免梯度累积。
  • 内存管理：大数据集使用 DataLoader 的 num_workers 和 pin_memory 优化加载。
  • 调试：使用 torch.set_printoptions() 控制张量输出格式，或 torch.autograd.set_detect_anomaly(True) 检测梯度问题。

6. 官方参考资源

  • PyTorch 官方文档：https://pytorch.org/docs/stable/index.html
  • torch 模块：https://pytorch.org/docs/stable/torch.html
  • torch.nn：https://pytorch.org/docs/stable/nn.html
  • torch.optim：https://pytorch.org/docs/stable/optim.html
  • torch.autograd：https://pytorch.org/docs/stable/autograd.html
  • 教程：https://pytorch.org/tutorials/
  • API 参考：https://pytorch.org/docs/stable/torch.html
  • GitHub 仓库：https://github.com/pytorch/pytorch
  • 社区论坛：https://discuss.pytorch.org/

7. 快速示例：完整训练流程

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset

# 数据
data = torch.randn(100, 10)
labels = torch.randint(0, 2, (100,))
dataset = TensorDataset(data, labels)
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

# 模型
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc = nn.Linear(10, 2)

    def forward(self, x):
        return self.fc(x)

model = Net()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 训练
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)
for epoch in range(5):
    for inputs, targets in dataloader:
        inputs, targets = inputs.to(device), targets.to(device)
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()
    print(f'Epoch {epoch+1}, Loss: {loss.item()}')

8. 进一步帮助

如果你需要更详细的 torch 模块说明（例如特定函数、优化器、或高级用法），或者有具体的代码调试需求，请提供更多细节，我可以为你提供定制化的解答或代码示例！

类似文章

发表回复

您的邮箱地址不会被公开。 必填项已用 * 标注