PyTorch张量操作大全从基础索引到高级广播机制的完整指南

张量基础：PyTorch的核心数据结构

PyTorch张量（Tensor）是该框架中最基本的数据结构，类似于NumPy的多维数组，但关键优势在于其能够在GPU上运行以加速计算，并支持自动微分。理解张量的基本操作是掌握深度学习的基础。所有深度学习模型本质上都是在操作张量——无论是输入数据、模型参数还是梯度信息，都以张量的形式存在。我们可以使用`torch.tensor()`函数或`torch.from_numpy()`从Python列表或NumPy数组创建张量，并使用`.dtype`、`.shape`和`.device`属性分别查看张量的数据类型、形状和存储设备（CPU或GPU）。

张量的创建与基本属性

创建张量的方法多样，常用的包括：`torch.zeros()`创建全零张量，`torch.ones()`创建全一张量，`torch.randn()`创建服从标准正态分布的随机张量，以及`torch.arange()`创建等差序列张量。这些函数都可以通过`dtype`参数指定数据类型（如`torch.float32`、`torch.int64`），通过`device`参数指定设备位置。张量的形状可以通过`.reshape()`、`.view()`或`.resize_()`方法进行调整，但需要注意这些方法在内存共享机制上的差异。

张量索引与切片：精准定位数据

PyTorch张量的索引和切片语法与Python列表和NumPy数组高度相似，这使得数据处理变得直观高效。基础索引使用`[ ]`操作符，可以指定单个索引值、切片范围或步长来访问张量的特定部分。例如，对于一个二维张量`tensor[i, j]`可以访问第i行第j列的元素，而`tensor[i1:i2, j1:j2]`则可以获取一个子矩阵。此外，布尔索引和整数数组索引提供了更灵活的数据选择方式，允许根据条件或特定索引列表筛选元素。

高级索引技巧

除了基本索引外，PyTorch还支持使用`torch.masked_select()`进行掩码选择，返回满足条件的一维张量；使用`torch.index_select()`沿指定维度选择索引对应的切片；以及使用`torch.gather()`根据索引张量从输入张量中收集特定元素。这些高级索引操作在处理不规则数据或实现复杂算法时尤为重要，例如在序列模型或强化学习中。

张量运算：数学操作与线性代数

PyTorch提供了丰富的数学运算函数，涵盖基本的算术运算、三角函数、指数对数函数以及完整的线性代数操作。算术运算如加法（`+`或`torch.add()`）、乘法（``或`torch.mul()`）等支持逐元素操作，也可以使用`@`操作符或`torch.matmul()`进行矩阵乘法。线性代数模块`torch.linalg`包含了更专业的操作，如矩阵分解（`torch.linalg.cholesky()`）、特征值计算（`torch.linalg.eig()`）和范数计算（`torch.linalg.norm()`），这些是构建复杂神经网络层和优化算法的基础。

归约操作与统计函数

归约操作是指对张量中所有元素或沿特定维度进行汇总计算的函数，例如`torch.sum()`求和、`torch.mean()`求平均值、`torch.max()`求最大值和`torch.argmax()`求最大值索引。这些函数通常接受`dim`参数指定计算维度，并可通过`keepdim=True`保持输出张量的维度数不变，这在后续的广播操作中非常有用。统计函数如`torch.std()`计算标准差、`torch.var()`计算方差，以及分位数计算`torch.quantile()`，为数据分析和模型评估提供了便利。

张量形状操作：重塑与组合

在处理深度学习数据时，经常需要改变张量的形状以适应模型输入或进行批量处理。`view()`和`reshape()`是最常用的形状重塑方法，可以将张量转换为指定形状，但`view()`要求张量在内存中连续存储，而`reshape()`会自动处理非连续情况。`squeeze()`和`unsqueeze()`分别用于删除和添加大小为1的维度，这在处理单样本数据或调整广播维度时特别有用。此外，`expand()`和`repeat()`都可以扩展张量的大小，但`expand()`不会复制数据，更节省内存。

张量的连接与分割

PyTorch提供了多种张量组合操作：`torch.cat()`沿现有维度连接多个张量，`torch.stack()`在新维度上堆叠张量，而`torch.chunk()`和`torch.split()`则用于将张量分割成多个部分。这些操作在构建神经网络时极其重要，例如将多个特征图拼接在一起，或将批量数据分割为小批量进行训练。

广播机制：智能维度扩展

广播是PyTorch中处理不同形状张量间运算的强大机制，其核心思想是自动扩展较小张量的维度，使其与较大张量的形状兼容，从而支持逐元素操作。广播遵循特定规则：首先从尾部维度开始对齐，比较两个张量的每个维度大小，如果维度大小相等，或其中一个大小为1，或其中一个张量在该维度上缺失，则可以进行广播。广播机制避免了显式复制数据的内存开销，使代码更简洁高效。

广播的实际应用示例

一个典型的广播例子是标量与张量的运算：当一个标量与多维张量相加时，标量会被广播到与张量相同的形状。同样，一个形状为(3,1)的列向量与形状为(1,4)的行向量相加，会通过广播得到形状为(3,4)的矩阵。理解广播规则对于调试张量形状不匹配错误至关重要，可以使用`torch.broadcast_tensors()`函数显式查看广播后的张量形状，帮助排查问题。

原地操作与内存管理

PyTorch中的大多数操作会创建新的张量，但有些操作提供了原地修改的版本，通过在函数名后添加下划线来表示（如`x.add_(y)`）。原地操作可以节省内存，避免创建临时张量，但在自动微分中需要谨慎使用，因为它们可能会破坏计算图，导致梯度计算错误。对于需要高效内存使用的场景，如处理大模型或受限设备，合理使用原地操作和内存共享机制（如`x.data`或`x.detach()`）可以显著优化性能。

梯度计算与自动微分

PyTorch的核心特性之一是自动微分，通过`autograd`模块实现。当张量的`requires_grad`属性设置为True时，PyTorch会跟踪在其上执行的所有操作，构建动态计算图。调用`.backward()`方法后，PyTorch会自动计算梯度并存储在张量的`.grad`属性中。理解张量与计算图的关系，以及如何使用`torch.no_grad()`上下文管理器禁用梯度跟踪，对于高效训练模型和进行推理至关重要。

高级张量操作：专业化应用

除了基础操作外，PyTorch还提供了众多高级张量操作，满足特定领域的专业需求。`torch.einsum()`函数支持爱因斯坦求和约定，可以简洁表达复杂的张量 contraction 操作。`torch.where()`实现条件选择，根据条件张量从两个张量中选择元素。散射和收集操作（`scatter_`和`gather`）在处理稀疏数据或特定索引模式时非常有用。此外，`torch.roll()`可以沿指定维度循环移动张量元素，常用于信号处理或数据增强。

与NumPy的互操作性

PyTorch张量与NumPy数组之间可以无缝转换，通过`tensor.numpy()`将张量转换为NumPy数组，或使用`torch.from_numpy()`将NumPy数组转换为张量。这种互操作性使得可以结合两个生态系统的优势，但需要注意共享内存的情况——在CPU上，这两种数据结构通常共享底层内存，修改一个会影响另一个，除非显式调用`copy()`方法。