很多人刚开始上手 PyTorch 时都会遇到一种奇妙的状态用nn.Linear、nn.Sequential搭模型跑 MNIST、CIFAR 都顺顺利利loss 也在降准确率也在涨。但只要一脱离“标准模板”想自己实现一个新的损失函数、给注意力机制加个 mask、或者手动拼接一下多模态输入特征代码立刻开始报维度错误甚至不报错但结果完全不对。这时候才意识到真正卡住你的不是对模型原理不懂而是对 PyTorch 最底层的张量运算缺乏“肌肉记忆”。本文是 PyTorch 系列教程的第 2 课核心就聚焦三件事逐元素计算、矩阵乘法、广播机制。先说判断这三类运算看起来简单实际构成了你在 PyTorch 中写的绝大多数代码的地基。不管是自定义模型、手写损失函数还是阅读开源库源码最后都会落到“这个张量该用*还是”“这个形状能不能广播”“矩阵乘法后 shape 是多少”这三个问题上。读完本文你能做到三件事彻底分清逐元素乘法*和矩阵乘法不再靠试错理解广播机制的真正规则能独立判断两个张量能不能做运算以及结果形状是什么不依赖nn.Linear用纯张量运算手写一个二分类前向计算和损失函数。文章里的代码都可以直接复制运行建议打开 Jupyter Notebook 或 VS Code 边看边敲遇到报错不要怕第 8 节专门整理了常见的维度问题排查表。1. 为什么张量运算值得单独用一节课讲很多 PyTorch 教程会把torch.Tensor的基本运算压缩在一节里讲完加法乘法就结束了。但实际工程里你会反复遇到下面这些场景场景一自定义损失函数。假设你在做目标检测需要计算 Smooth L1 Loss或者在做对比学习需要自己算 InfoNCE Loss。这些损失函数看起来只有几行公式但落到代码里全是张量之间的逐元素计算、矩阵乘法、以及高维张量的 reshape 和广播。如果你对*和的区别不敏感一个符号写错模型可能不报错但训练出的结果就是错的而且这种错误非常难排查。场景二处理批量数据。深度学习中几乎不会一次只处理一个样本而是把一批样本叠成一个张量形状通常是(batch_size, ...)。这个时候线性层的前向计算就变成x w.T b其中权重和偏置要借助广播机制自动适配 batch 维。如果不懂广播你会发现维度怎么凑都别扭。场景三阅读开源模型源码。Transformer 里的自注意力公式是softmax(Q K.transpose(-2, -1) / sqrt(d_k)) V。任何一个从事 NLP 的开发者都会经常与这段代码打交道它本质就是一个高维矩阵乘法和逐元素的缩放操作。你不需要一行行背代码但必须能在心里模拟每个步骤的 shape 变化。这恰恰是很多人跳过的一环。框架封装得越好底层运算越容易被忽略可一旦遇到问题又必须回到底层才能解释清楚。我在教学和实际项目中看到的规律是能熟练说出每个中间张量shape的开发者写代码的出错率会明显低于只知道“调 API”的人。因为 shape 一旦能在头脑中推演PyTorch 的报错信息对你来说就不再是乱码而是很明确的提示。所以这一课虽然“基础”但它决定了你后续阅读源码、写自定义操作、做模型调试能不能顺畅。下面我们从张量本身开始一步一步把这套规则建立起来。2. 张量基础维度、形状与数据类型2.1 张量到底是什么张量Tensor可以理解为“可以多维排列的数值数组”。0 维张量是标量1 维张量是向量2 维张量是矩阵3 维及以上的张量在深度学习里非常常见。例如一批 28×28 的灰度图片形状是(batch_size, 1, 28, 28)这是 4 维张量一批文本序列经过 embedding 之后形状是(batch_size, seq_len, hidden_size)这是 3 维张量一个全连接层的权重矩阵形状是(in_features, out_features)这是 2 维张量。在 PyTorch 中一切数据都是张量。神经网络的前向传播本质上就是张量依次经过一系列运算形状不断发生变换的过程。2.2 张量的基本属性每个张量都有几个关键属性调试时几乎天天用属性含义示例shape张量每个维度的大小torch.Size([100, 3])ndim张量维度数量也叫秩二维矩阵的ndim2dtype数据类型torch.float32device数据存放在 CPU 还是 GPUcpu或cuda:0其中dtype和device是很多隐蔽 bug 的来源。两个张量即使形状完全一致只要一个是torch.float32、另一个是torch.float64在很多运算中就会报错或不满足预期。GPU 张量和 CPU 张量也不能直接混合运算。这一点在后面的实践中会反复提到。2.3 常用创建方式为了熟练操作我们先从一个最小示例开始创建一个张量并查看它的属性。import torch # 从 Python 列表创建 a torch.tensor([[1, 2, 3], [4, 5, 6]]) # 创建全 0 张量 b torch.zeros(2, 3) # 创建全 1 张量 c torch.ones(2, 3) # 创建随机张量标准正态分布 d torch.randn(2, 3) # 从 0 到 11 的序列并 reshape 成 3 行 4 列 e torch.arange(12).reshape(3, 4) print(a.shape:, a.shape) print(a.ndim:, a.ndim) print(a.dtype:, a.dtype) print(b.device:, b.device)运行结果大致如下a.shape: torch.Size([2, 3]) a.ndim: 2 a.dtype: torch.int64 b.device: cpu注意没有指定dtype时torch.tensor会从 Python 整数推算出torch.int64而torch.zeros、torch.randn默认创建的是torch.float32。深度学习中绝大多数计算都要求浮点数所以建模时通常是先创建浮点张量或者在.float()转换后再运算。张量创建本身不是本文重点但它是后面所有运算的载体。建议你先多创建几个不同 shape 的张量打印出来看看“长什么样”熟悉了之后再进入运算规则。3. 环境准备安装 PyTorch 并验证可用性如果你还没有装 PyTorch这里用最简方式带你过一遍。安装 Python 后直接用 pip 安装即可。对于只打算学习基础运算、先跑通代码的读者CPU 版本就够用了。本文所有示例在 CPU 上都能运行不需要 GPU。# CPU 版本适合学习和调试 pip install torch如果你使用 GPU 并希望利用 CUDA 加速建议访问 PyTorch 官网根据操作系统和 CUDA 版本复制安装命令。官网给出的命令一般长这样# 请将 cuda_version 替换为官网列出的真实标识比如 cu118、cu121 等 pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cuda_version这里的cuda_version不应该凭记忆填写而是先通过nvidia-smi查看本机 CUDA 版本再对照官网推荐。选错版本最常见的现象是PyTorch 装成功了但torch.cuda.is_available()返回False。安装完成后在 Python 交互环境或脚本里执行import torch print(torch.__version__) print(torch.cuda.is_available()) if torch.cuda.is_available(): print(torch.cuda.get_device_name(0))如果打印出了版本号说明安装成功。torch.cuda.is_available()为True则表示 GPU 可用。对于本节的学习即使它是False也不影响你运行任何代码。另外很多读者会遇到下载速度慢的问题这通常和网络环境有关可以配置国内 PyPI 镜像加速。具体配置方式以你使用的镜像源官方说明为准。环境准备到这里就足够了我们接着进入主题。4. 逐元素计算加减乘除、比较与取余4.1 逐元素计算的含义逐元素计算Element-wise Operation是指两个张量对应位置上的元素各自完成运算。例如两个形状相同的张量相加结果是同一位置的数字相加得到一个新张量。这种运算在深度学习中极其常见给数据做归一化、在损失函数中相减、对预测结果的每个元素取平方、比较预测值和真实值是否一致……本质上都是逐元素计算。4.2 运算符与函数对照表PyTorch 为每种逐元素运算都提供了两种写法一种是直接的 Python 运算符另一种是torch下的函数。两者效果一致用哪种看代码风格和个人习惯。运算运算符写法函数写法加法a btorch.add(a, b)减法a - btorch.sub(a, b)逐元素乘法a * btorch.mul(a, b)除法a / btorch.div(a, b)幂运算a ** 2torch.pow(a, 2)取余a % 2torch.remainder(a, 2)整除a // 2torch.floor_divide(a, 2)大于a 1torch.gt(a, 1)小于a 1torch.lt(a, 1)等于a 1torch.eq(a, 1)来看一个完整示例import torch x torch.tensor([1.0, 2.0, 3.0]) y torch.tensor([4.0, 5.0, 6.0]) print(x y , x y) print(x - y , x - y) print(x * y , x * y) print(x / y , x / y) print(x ** 2 , x ** 2) print(x % 2 , x % 2) print(x 1 , x 1)运行结果x y tensor([5., 7., 9.]) x - y tensor([-3., -3., -3.]) x * y tensor([ 4., 10., 18.]) x / y tensor([0.2500, 0.4000, 0.5000]) x ** 2 tensor([1., 4., 9.]) x % 2 tensor([1., 0., 1.]) x 1 tensor([False, True, True])比较运算的结果是布尔张量在后续项目中常用作掩码mask。比如要从一批张量里筛选出大于某个阈值的元素就可以直接用布尔索引x torch.tensor([1.0, 2.0, 3.0, 4.0]) mask x 2 print(mask) print(x[mask])这里x[mask]会把所有满足条件的元素取出来得到tensor([3., 4.])。在 Transformer 的注意力掩码、目标检测的置信度过滤等场景里这种模式非常常见。4.3 最容易混淆的*与初学者最容易混淆的就是*和因为数学里我们习惯把矩阵相乘写成乘号。*表示逐元素乘法element-wise multiplication对应torch.mul表示矩阵乘法matrix multiplication对应torch.matmul。用一个简单的例子说明区别import torch a torch.tensor([[1, 2], [3, 4]]) b torch.tensor([[5, 6], [7, 8]]) print(逐元素乘法 a * b:) print(a * b) print(矩阵乘法 a b:) print(a b)输出逐元素乘法 a * b: tensor([[ 5, 12], [21, 32]]) 矩阵乘法 a b: tensor([[19, 22], [43, 50]])从输出可以清晰看到*是把相同位置的元素相乘则是标准的矩阵乘法结果里的每个元素是左侧行向量与右侧列向量的内积。在深度学习中*常用于逐元素缩放、掩码相乘、损失函数里的逐项计算常用于全连接层、注意力矩阵、投影运算。写代码之前先问一句“我现在是做逐元素操作还是做矩阵变换”能帮你避免一大类隐蔽错误。5. 矩阵乘法mm、bmm、matmul 与 运算符5.1 矩阵乘法的维度约束矩阵乘法的数学定义是如果A的形状是(m, k)B的形状是(k, n)那么A B的结果形状是(m, n)。这里的关键约束是A 的最后一维必须等于 B 的倒数第二维也就是两个矩阵的内侧维度要相等。至于外侧的两个维度m和n没有任何限制。很多人写矩阵乘法报错都是因为这个约束。错误信息通常会提示RuntimeError: size mismatch, got input (m, k) and (k, n), where k ! k解决方法是打印出两个张量的shape确认内侧维度是否一致。如果维度反了可以用.T或.transpose()转置后再乘。5.2 三种常用 API 如何选PyTorch 里有三个经常见到的矩阵乘法 API很多人会混淆它们的使用场景。API适用场景注意事项torch.mm只处理二维矩阵输入必须是两个二维张量高维会报错torch.matmul/支持高维支持广播最常用推荐优先使用torch.bmm批量矩阵乘法输入必须是两个三维张量batch 维必须一致torch.mm是最古老的接口现在新的代码里已经很少单写了。是torch.matmul的运算符写法功能完全一致代码更简洁。torch.bmm在批量场景中很常用。假设你有一个 batch 的查询矩阵q形状是(batch_size, seq_len, head_dim)还有一个键矩阵k形状是(batch_size, seq_len, head_dim)。你想对每个 batch 分别做矩阵乘法q k.transpose(-2, -1)就能完成。如果把它们扩展成三维使用torch.bmm也等价。来看一个torch.bmm的示例import torch # 4 个 batch每个是 3x5 的矩阵 a torch.randn(4, 3, 5) # 4 个 batch每个是 5x2 的矩阵 b torch.randn(4, 5, 2) c torch.bmm(a, b) print(c.shape) # torch.Size([4, 3, 2])结果说明batch 维保持不变后面两个维度的矩阵乘法结果就是(3, 2)。5.3 一个手写线性层的例子nn.Linear的内部本质就是一次矩阵乘法加广播加法。我们用纯张量运算来实现它让自己对 shape 有更直观的感受import torch # 模拟 100 个样本每个样本有 3 个特征 x torch.randn(100, 3) # 随机初始化权重输入特征 3输出特征 1 w torch.randn(3, 1) # 偏置形状为 (1,) b torch.randn(1) # 前向计算y x w b y x w b print(x.shape:, x.shape) print(w.shape:, w.shape) print(y.shape:, y.shape)这里发生了什么x形状是(100, 3)w形状是(3, 1)矩阵乘法结果是(100, 1)偏置b形状是(1,)在加法时自动广播成(100, 1)最终y的形状是(100, 1)代表 100 个样本各自的线性输出。这个例子同时也展示了广播机制的作用。下面一节我们把广播讲透。6. 广播机制自动扩展维度的三条规则与常见坑6.1 什么是广播广播Broadcasting指两个形状不完全相同的张量做运算时PyTorch 自动扩展其中一个张量的维度使两者形状匹配然后进行逐元素运算。它的价值是简化代码。没有广播我们要手动把b复制成 100 行才能和x w相加。有了广播PyTorch 自动处理扩展逻辑既省内存又省代码。但是广播不等于“两个张量随便算”。PyTorch 有一套严格的规则不符合规则时直接报错。6.2 三条规则广播规则可以浓缩成三条规则一从最后一个维度开始向前的对齐。两个张量比较形状时先看各自的最后一个维度然后再向前依次比较。规则二每个维度上要么大小相同要么其中一个为 1。如果两个张量在某个维度上的大小相等那这个维度直接匹配如果其中一个为 1就可以拉伸到另一个的大小如果两者都不为 1 且不相等广播失败。规则三缺失维度按 1 处理。如果一个张量的维度数量比另一个少缺失的维度会被当作 1 来参与对齐。这三条规则合起来可以用一句话记忆从右往左逐维对齐相等或一方为 1 就能广播。6.3 典型例子把规则落到代码里最有代表性的一种情况是import torch # (3, 1) 张量 a torch.tensor([[1], [2], [3]]) # (1, 4) 张量 b torch.tensor([[10, 20, 30, 40]]) # 广播相加结果应该是 (3, 4) c a b print(a.shape:, a.shape) print(b.shape:, b.shape) print(c.shape:, c.shape) print(c)输出a.shape: torch.Size([3, 1]) b.shape: torch.Size([1, 4]) c.shape: torch.Size([3, 4]) tensor([[11, 21, 31, 41], [12, 22, 32, 42], [13, 23, 33, 43]])这个过程可以理解为a的每一行被扩展了 4 次b的每一行被扩展了 3 次最后逐元素相加。更常见的场景是向量与矩阵相加。例如对一批数据做 z-score 归一化需要把每个样本减去均值向量import torch # 3 个样本4 个特征 x torch.randn(3, 4) # 每个特征的均值 mean x.mean(dim0) # 每个特征的标准差 std x.std(dim0) # 广播x 是 (3, 4)mean 是 (4,)缺失维度按 1 处理所以等效于 (1, 4) x_normalized (x - mean) / std print(x.shape:, x.shape) print(mean.shape:, mean.shape) print(x_normalized.shape:, x_normalized.shape)这里mean和std的形状是(4,)与(3, 4)的x做运算时缺失的维度被当作 1结果正常广播成(3, 4)。6.4 显式扩展expand 与 broadcast_to广播是隐式的代码里没有显式的扩展过程。如果希望把某个张量显式扩展成目标形状可以使用expand或broadcast_to。import torch a torch.tensor([[1], [2], [3]]) # (3, 1) # 显式扩展成 (3, 4) b a.expand(3, 4) print(expand 结果 shape:, b.shape) # 也可以直接广播到目标形状 c a.broadcast_to(3, 4) print(broadcast_to 结果 shape:, c.shape)注意expand返回的新张量仍然共享底层数据并没有真的复制一份因此它不会占用额外内存。这个特性在实现大规模模型时很重要。6.5 广播失败和容易误用的场景广播虽然方便但也会带来隐蔽问题。常见的情况是两个张量本来形状不同你预期它们不能直接相加但广播机制“好心”地把它们拉在一起导致结果形状和你以为的不一致而且不报错。看这个例子import torch # 3 个样本4 个特征 x torch.randn(3, 4) # 这是一个形状为 (4, 3) 的矩阵 w torch.randn(4, 3) # 想给 x 的每个样本加一个向量但这个向量的形状不是 (4,)而是 (3,) v torch.randn(3) # 结果形状是 (4, 3)而不是 (3, 4) result w v print(w.shape:, w.shape) print(v.shape:, v.shape) print(result.shape:, result.shape)w的形状是(4, 3)v的形状是(3,)从右往左对齐时w的最后一维是 3v的最后一维也是 3正好匹配所以v会被当作一行广播到每一行。这种运算不会报错但如果你原本想的是把v当作一列加到每个特征的维度上结果就和预期完全不同。防止这种问题的最佳实践是在关键运算前显式打印shape或者用assert校验维度。第 9 节会专门给出工程建议。广播机制的另一个教训是不要以为“能运行”就是正确的。PyTorch 的广播规则是固定且严格的它能帮你省去很多手动 reshape 的麻烦但也可能掩盖逻辑错误。只有在理解规则后主动检查 shape才能写出既简洁又可靠的前向计算代码。7. 综合实战用三类张量运算实现一个小模型的前向计算现在把三类运算综合到一起写一个不依赖nn.Linear的二分类前向计算和损失函数。这个例子能覆盖矩阵乘法x w广播机制偏置b的广播逐元素计算sigmoid 激活、损失函数中的逐项计算聚合运算mean()把损失聚合成标量。完整代码import torch # 1. 构造数据集100 个样本3 个特征 x torch.randn(100, 3) # 随机生成二分类标签0 或 1 y (torch.randn(100) 0).float() # 2. 初始化参数并设置需要梯度 w torch.randn(3, requires_gradTrue) b torch.randn(1, requires_gradTrue) # 3. 前向计算logits x w b # x 形状 (100, 3)w 形状 (3,)matmul 结果形状 (100,) # b 形状 (1,)广播到 (100,) logits x w b # 4. 逐元素计算 sigmoidp 1 / (1 exp(-logits)) p 1 / (1 torch.exp(-logits)) # 5. 二分类交叉熵损失手写形式 eps 1e-8 loss -(y * torch.log(p eps) (1 - y) * torch.log(1 - p eps)).mean() # 6. 反向传播 loss.backward() print(logits.shape:, logits.shape) print(p.shape:, p.shape) print(p 的范围: [{:.4f}, {:.4f}].format(p.min(), p.max())) print(loss:, loss.item()) print(w.grad.shape:, w.grad.shape) print(b.grad.shape:, b.grad.shape)这段代码的关键点x w中x是(100, 3)w是(3,)。PyTorch 的matmul对一维张量有特殊规则把w当作(3, 1)计算后把最后一维去掉所以结果是(100,)。logits b中logits是(100,)b是(1,)广播后相加结果仍是(100,)。sigmoid 是对logits的每个元素做逐元素计算。损失函数里的y * torch.log(p)和(1 - y) * torch.log(1 - p)都是逐元素乘法最后用mean()把所有元素平均成一个标量。loss.backward()会沿着计算图自动求导这也是下一课autograd的重点。有了w.grad和b.grad就可以继续手动实现一次简单的梯度下降更新这也是你理解 PyTorch 优化器工作原理的第一步。8. 运行结果、效果验证与常见问题排查8.1 预期输出运行上面的综合实战代码你会看到类似下面的输出logits.shape: torch.Size([100]) p.shape: torch.Size([100]) p 的范围: [0.2345, 0.8765] loss: 0.6821 w.grad.shape: torch.Size([3]) b.grad.shape: torch.Size([1])由于随机种子不同具体数值会和这里不完全一致但有几个判断标准是确定的logits.shape必须是torch.Size([100])说明矩阵乘法和广播正确p的每个元素都必须在 0 到 1 之间loss是一个标量0 维张量打印时通过.item()拿到 Python 浮点数w.grad和b.grad不为None说明反向传播成功。如果运行报错优先检查是不是x w这一步维度不匹配。把x改成(100, 3)、w改成(3,)就能复现文章里的场景。8.2 常见问题与排查思路下面这张表整理了初学者在张量运算中最常遇到的问题问题现象可能原因排查方式解决方案RuntimeError: size mismatch矩阵乘法内侧维度不一致打印两个张量的shape确认a.shape[-1] b.shape[-2]转置或调整维度RuntimeError: The size of tensor a ... must match ...广播失败两个张量在某维度都不为 1 且不相等对照广播规则从右往左逐维检查用reshape、unsqueeze补充维度程序运行成功但结果数值不对混淆了*与检查运算意图逐元素还是矩阵变换逐元素用*矩阵乘法用torch.bmm报错bmm要求两个输入都是三维且 batch 维一致打印输入shape改用/matmul或先reshapeGPU 版 PyTorch 装上后torch.cuda.is_available()为 FalseCUDA 版本与 PyTorch 不匹配运行nvidia-smi查看驱动支持的 CUDA 版本按官网命令重新安装对应版本前向传播正常但loss一直是nanlog(0)、除零或梯度爆炸打印p的范围检查p是否出现 0 或 1加eps或使用数值稳定的损失函数实现排查维度问题有一个通用套路不要猜直接在报错行之前加print(x.shape)把中间所有张量的 shape 打印出来。大多数维度错误三步之内就能定位。9. 最佳实践与下一步学习方向9.1 五个值得养成的好习惯第一在关键位置打印 shape。调试模型时把每个中间张量的 shape 打印出来比读报错日志快得多。尤其是自定义模块、自定义损失函数运行前先用一小批假数据做前向验证。第二善用 assert 做维度约束。如果某个层要求输入必须是(batch, seq_len, hidden)可以在实现里加一句断言assert x.dim() 3, f期望输入是 3 维实际得到 {x.shape}这样上线时维度一旦变化程序会立刻以显眼方式报错而不是在后续运算中产生莫名其妙的数值错误。第三写清楚“维度配方”注释。在实现模型时用注释标注每个张量的形状变化。例如# x: (batch, seq_len, hidden) # w: (hidden, out_features) # out: (batch, seq_len, out_features) out x w这种注释成本极低但能极大提高代码可读性也更方便团队协作。第四关注 dtype 和 device。不要只盯着 shape。两个形状完全一致的张量一个float32、一个float64在某些操作中也会出现意想不到的问题。统一在模型入口做tensor.float()把数据迁移到同一设备能避免很多隐性问题。第五谨慎使用 inplace 运算。例如x.add_(1)会直接修改原张量。在 autograd 机制下inplace 运算有时会破坏计算图导致梯度计算报错。除非你明确知道自己在做什么否则优先使用非 inplace 写法比如x x 1。9.2 下一步学什么掌握了张量运算规则之后建议按