Skip to main content

Pytorch

PyTorch 张量 (Tensors)

如果你已经熟悉 NumPy 的 ndarray，那么理解 PyTorch 的张量会非常容易。PyTorch 张量与 NumPy 数组非常相似，但有一个关键的区别：张量可以在 GPU 上进行运算，从而极大地加速计算，这对于训练大模型至关重要。

  • 什么是张量？

    • 多维数组。
    • 0维张量：标量 (一个数字)
    • 1维张量：向量 (一列数字)
    • 2维张量：矩阵
    • 更高维度的张量... (例如，在 LLM 中，一个批次的词嵌入数据通常是 3D 张量：[batch_size, sequence_length, embedding_dimension])

  • 创建张量:

    import torch
import numpy as np

# 从 Python 列表创建
data_list = [[1, 2], [3, 4]]
tensor_from_list = torch.tensor(data_list)
print("从列表创建的张量:\n", tensor_from_list)
print("张量的数据类型:", tensor_from_list.dtype) # 默认为 torch.int64

# 从 NumPy 数组创建 (共享内存，除非显式拷贝)
numpy_array = np.array([[5, 6], [7, 8]])
tensor_from_numpy = torch.from_numpy(numpy_array) # 或者 torch.as_tensor(numpy_array)
print("\n从NumPy数组创建的张量:\n", tensor_from_numpy)

# NumPy 数组和 PyTorch 张量之间的转换
numpy_back = tensor_from_numpy.numpy()
print("\n转换回NumPy数组:\n", numpy_back)

# 创建特定形状和类型的张量
zeros_tensor = torch.zeros(2, 3) # 2x3 的全零张量
ones_tensor = torch.ones(2, 3, dtype=torch.float32) # 指定数据类型
rand_tensor = torch.rand(2, 3)  # [0, 1) 均匀分布
randn_tensor = torch.randn(2, 3) # 标准正态分布

print("\n全零张量:\n", zeros_tensor)
print("全一张量 (float32):\n", ones_tensor)
print("随机张量 (均匀分布):\n", rand_tensor)
print("随机张量 (正态分布):\n", randn_tensor)

# 指定数据创建张量，并让PyTorch推断大小
x_data = torch.tensor([[1., -1.], [1., -1.]]) # 注意使用浮点数
x_ones = torch.ones_like(x_data) # 创建形状和类型与x_data相同的全一张量
x_rand = torch.rand_like(x_data, dtype=torch.float) # 覆盖数据类型
print("\nx_rand:\n", x_rand)

  • 张量属性: 与 NumPy 类似

    • tensor.shape 或 tensor.size(): 张量的形状。
    • tensor.dtype: 张量的数据类型 (如 torch.float32, torch.int64)。
    • tensor.device: 张量所在的设备 (CPU 或 GPU)。
    print(f"\n形状: {rand_tensor.shape}")
print(f"数据类型: {rand_tensor.dtype}")
print(f"设备: {rand_tensor.device}") # 默认在 CPU 上

  • 张量操作: 许多操作与 NumPy 类似。

    • 算术运算: +, -, *, /, @ (矩阵乘法) 或 torch.matmul()
    • 索引和切片: 与 NumPy 类似。
    • 变形: tensor.view(), tensor.reshape(), tensor.transpose(), tensor.permute()
      • view(): 返回一个新的张量，与原始张量共享数据。只能用于连续内存的张量。
      • reshape(): 功能更强大，不总是共享数据，可以在需要时创建副本。
    • 聚合操作: torch.sum(), torch.mean(), torch.max(), torch.argmax()
    a = torch.tensor([[1, 2], [3, 4]], dtype=torch.float32)
b = torch.tensor([[5, 6], [7, 8]], dtype=torch.float32)

# 加法
print("\na + b:\n", a + b)
print("torch.add(a, b):\n", torch.add(a, b))

# 矩阵乘法
print("\na @ b (矩阵乘法):\n", a @ b)
print("torch.matmul(a, b):\n", torch.matmul(a, b))

# 索引
print("\na的第一行:", a[0])
print("a的[0,1]元素:", a[0, 1])

# view (类似 reshape, 但共享数据)
x = torch.randn(4, 4)
y = x.view(16)
z = x.view(-1, 8)  # -1 表示由其他维度推断
print("\nx的形状:", x.shape)
print("y的形状 (view成16个元素):", y.shape)
print("z的形状 (view成nx8):", z.shape)

GPU 加速

这是 PyTorch 相对于 NumPy 的一个核心优势，对于大模型训练至关重要。

  • 检查 GPU 是否可用:

    if torch.cuda.is_available():
    device = torch.device("cuda")          # 使用第一个可用的 CUDA GPU
    print(f"\nGPU '{torch.cuda.get_device_name(0)}' 可用，将使用GPU。")
else:
    device = torch.device("cpu")
    print("\nGPU 不可用，将使用CPU。")

  • 将张量移动到 GPU:

    # 创建一个张量 (默认在CPU)
cpu_tensor = torch.randn(3, 3)
print("CPU 张量设备:", cpu_tensor.device)

# 将张量移动到之前确定的设备 (GPU或CPU)
if torch.cuda.is_available():
    gpu_tensor = cpu_tensor.to(device)
    print("GPU 张量设备:", gpu_tensor.device)

    # 也可以直接在GPU上创建张量
    gpu_direct_tensor = torch.randn(2, 2, device=device)
    print("直接在GPU上创建的张量设备:", gpu_direct_tensor.device)

    # 注意：不同设备上的张量不能直接运算，需要先移动到同一设备
    # cpu_tensor + gpu_tensor # 这会报错
    result_on_gpu = gpu_tensor + gpu_direct_tensor
    print("GPU上运算结果设备:", result_on_gpu.device)

    # 将结果移回CPU (例如用于打印或与NumPy交互)
    result_on_cpu = result_on_gpu.cpu()
    print("移回CPU的结果设备:", result_on_cpu.device)

  • .to(device) 方法 是将张量或模型移动到特定设备（CPU 或 GPU）的通用方法。

自动求导 (torch.autograd)

这是 PyTorch 另一个核心特性，也是训练神经网络（包括 LLM）的基石。它能够自动计算梯度（导数）。

  • requires_grad=True: 当创建一个张量时，如果设置 requires_grad=True，PyTorch 会开始追踪在该张量上的所有操作，以便进行梯度计算。
  • 计算图 (Computational Graph): PyTorch 会构建一个动态计算图，记录数据如何通过操作组合得到最终结果。
  • .backward(): 当你在一个标量输出上调用 .backward() 时，PyTorch 会自动计算计算图中所有 requires_grad=True 的张量相对于该标量输出的梯度。
  • .grad: 梯度会累积到张量的 .grad 属性中。
# 创建需要梯度的张量
x = torch.ones(2, 2, requires_grad=True)
print("\nx:\n", x)

# 对 x 进行一些操作
y = x + 2
print("y:\n", y) # y 也会自动具有 grad_fn，因为它是由 x 计算得来的

z = y * y * 3
out = z.mean() # out 是一个标量
print("z:\n", z)
print("out (标量):\n", out)

# 计算梯度
out.backward() # 等价于 out.backward(torch.tensor(1.))

# 打印 x 的梯度 d(out)/dx
print("\nx的梯度 (d(out)/dx):\n", x.grad)

解释 x.grad 的计算过程 (链式法则): out = (1/4) * Σ z_i z_i = 3 * y_i^2 y_i = x_i + 2

∂out/∂z_i = 1/4 ∂z_i/∂y_i = 6 * y_i ∂y_i/∂x_i = 1

所以，∂out/∂x_i = (∂out/∂z_i) * (∂z_i/∂y_i) * (∂y_i/∂x_i) = (1/4) * (6 * y_i) * 1 = (3/2) * y_i 因为 x = [[1,1],[1,1]], 所以 y = [[3,3],[3,3]] 因此，x.grad 应该是 (3/2) * [[3,3],[3,3]] = [[4.5, 4.5],[4.5, 4.5]]，这与程序输出一致。

LLM 相关性:

  • 大模型的参数（权重和偏置）都是 requires_grad=True 的张量。
  • 损失函数 (Loss Function) 的输出是一个标量。
  • 调用 loss.backward() 会计算损失相对于模型所有参数的梯度。
  • 优化器 (Optimizer) 使用这些梯度来更新模型参数，从而使损失减小。

重要注意事项:

  • 只有浮点类型的张量 (如 torch.float32, torch.float64) 才能计算梯度。
  • 默认情况下，新创建的张量 requires_grad=False。
  • 梯度是累积的！在每次迭代（例如，在训练循环中）计算梯度之前，通常需要使用 optimizer.zero_grad() 或手动将 .grad 设置为 None 来清除旧的梯度。 
    # 梯度累积示例
# 第一次 backward
# out.backward()
# print(x.grad)

# 如果再次 backward 而不清除梯度
# out.backward() # 这会报错，因为图已经被释放了，除非指定 retain_graph=True
                 # 或者梯度会累加

# 正确的做法是在新一轮计算前清除梯度
# x.grad.data.zero_() # 一种清除梯度的方式
  • 可以用 with torch.no_grad(): 上下文管理器来临时禁用梯度计算，这在模型评估（推理）阶段非常有用，可以减少内存消耗并加速计算。 
    print("\n原始x的requires_grad:", x.requires_grad)
print("在torch.no_grad()上下文中:")
with torch.no_grad():
    y_no_grad = x + 2
    print("y_no_grad的requires_grad:", y_no_grad.requires_grad) # False

总结一下，主要包括：

  1. 张量的创建:

    • 从 Python 列表或 NumPy 数组创建 (torch.tensor(), torch.from_numpy())。
    • 创建特定形状和值的张量 (torch.zeros(), torch.ones(), torch.rand(), torch.randn(), torch.arange(), torch.linspace())。
    • 使用 _like 版本根据现有张量创建 (torch.zeros_like())。

  2. 张量的属性:

    • 形状 (.shape, .size())。
    • 数据类型 (.dtype)。
    • 所在设备 (.device)。
    • 是否需要梯度 (.requires_grad)。
    • 梯度值 (.grad)。
    • 梯度函数 (.grad_fn)，用于追踪操作历史以进行反向传播。

  3. 张量的基本运算:

    • 算术运算: 加、减、乘、除、矩阵乘法 (@ 或 torch.matmul())。
    • 元素级运算: torch.abs(), torch.sqrt(), torch.exp(), torch.log(), torch.sin(), torch.sigmoid(), torch.relu() 等。
    • 比较运算: torch.eq(), torch.gt(), torch.lt() 等，返回布尔张量。

  4. 索引、切片和连接:

    • 索引和切片: 与 NumPy 非常相似，用于访问和修改张量的部分元素。
    • 连接: torch.cat() (沿指定维度拼接张量), torch.stack() (沿新维度堆叠张量)。
    • 拆分: torch.split(), torch.chunk()。

  5. 形状变换:

    • tensor.view(): 改变张量形状（共享数据，要求内存连续）。
    • tensor.reshape(): 改变张量形状（不一定共享数据，更灵活）。
    • tensor.T 或 tensor.transpose(): 转置。
    • tensor.permute(): 任意维度重排。
    • tensor.unsqueeze(): 增加维度。
    • tensor.squeeze(): 移除大小为1的维度。

  6. 聚合操作:

    • torch.sum(), torch.mean(), torch.prod(), torch.std(), torch.var()。
    • torch.min(), torch.max() (返回数值)。
    • torch.argmin(), torch.argmax() (返回索引)。
    • 可以指定 dim 参数沿特定维度进行聚合。

  7. GPU 加速:

    • 检查 GPU 可用性 (torch.cuda.is_available())。
    • 将张量移至设备 (tensor.to(device))。
    • 直接在特定设备创建张量 (torch.tensor(..., device=device))。

  8. 自动求导 (torch.autograd):

    • 设置 requires_grad=True 来追踪操作。
    • 调用 .backward() 计算梯度。
    • 通过 .grad 属性访问梯度。
    • 使用 with torch.no_grad(): 禁用梯度计算。
    • 梯度累积和清除 (optimizer.zero_grad() 或 tensor.grad.zero_())。

这些操作构成了 PyTorch 张量操作的基石。当你开始构建神经网络层、定义损失函数和编写训练循环时，你会不断地用到这些基础操作。

当然，PyTorch 作为一个庞大的库，还有更多高级或特定场景下的张量操作，但对于入门和理解核心原理，以上这些是最重要的。 随着你学习的深入，比如接触到更复杂的网络结构（如 RNN、Transformer 中的特定操作）或者分布式训练，你可能会遇到一些新的、更专门的张量函数。但万变不离其宗，都是基于这些基础概念的扩展。

对于初步学习，掌握这些已经非常好了。下一步的关键是将这些孤立的张量操作与神经网络的概念（如层、激活函数、损失函数、优化器）联系起来。

Back to top