Skip to content

第 4 章 张量运算与广播

学习目标

  • 掌握逐元素算术运算与比较运算
  • 掌握原地运算 add_ 等及其注意点
  • 掌握矩阵乘法 @ / matmul,区分 *@
  • 掌握广播规则,能判断两个形状能否运算
  • 会用 torch.dot 做向量点积

4.1 逐元素运算

逐元素运算(elementwise operation)是形状相同的两个张量按位置一一对应做运算:

python
import torch

a = torch.tensor([1, 2, 3])
b = torch.tensor([4, 5, 6])
print(a + b)
print(a * b)
print(a - b)

输出:

tensor([5, 7, 9])
tensor([ 4, 10, 18])
tensor([-3, -3, -3])

算术运算符与 Python 一致,行为如下:

运算结果说明
a + 2tensor([3, 4, 5])加法
a / 2tensor([0.5000, 1.0000, 1.5000])除法返回浮点
a // 2tensor([0, 1, 1])整除,向下取整
a % 2tensor([1, 0, 1])取余
a ** 2tensor([1, 4, 9])
python
import torch

a = torch.tensor([1, 2, 3])
print(a / 2)
print(a // 2)
print(a % 2)
print(a ** 2)

输出:

tensor([0.5000, 1.0000, 1.5000])
tensor([0, 1, 1])
tensor([1, 0, 1])
tensor([1, 4, 9])

注意整数除法:1 除以 2 得到 0.5,但整除 1 // 2 得到 0——和 Python 的语义一致。

4.2 比较运算与布尔张量

比较运算逐元素返回 True / False,结果是布尔张量:

python
import torch

a = torch.tensor([1, 2, 3])
print(a > 2)
print((a > 1) & (a < 3))    # 与运算:两个条件都要满足

输出:

tensor([False, False,  True])
tensor([False,  True, False])

布尔张量是第 3 章布尔索引的基础。& 是与、| 是或、~ 是取反;多个比较条件用括号包起来再连接,否则会因运算符优先级出错。

4.3 原地运算

a + 1 返回一个新张量,不改动 a。以 _ 结尾的方法(如 add_)是原地运算(in-place),直接修改当前张量:

python
import torch

x = torch.tensor([1, 2, 3])
x.add_(10)
print(x)

输出:

tensor([11, 12, 13])

常见原地方法:add_mul_sub_div_。原地运算省内存,但对需要求梯度的张量有风险——PyTorch 会拦截对叶子张量的原地修改并报错(第 6 章详细讲)。没有明确理由时,优先用返回新张量的写法。

4.4 矩阵乘法:*@ 的区别

* 是逐元素相乘(形状必须兼容),@ 是矩阵乘法(遵循线性代数规则):

python
import torch

m1 = torch.tensor([[1., 2.],
                   [3., 4.]])
m2 = torch.tensor([[5., 6.],
                   [7., 8.]])
print(m1 @ m2)

输出:

tensor([[19., 22.],
        [43., 50.]])

验证第一个元素:1×5 + 2×7 = 19,这是标准的矩阵乘法。如果改成 m1 * m2,得到的是对应位置相乘 [[5, 12], [21, 32]]

向量与矩阵相乘时,向量当作行或列自动适配:

python
import torch

v = torch.tensor([1., 2.])
m1 = torch.tensor([[1., 2.],
                   [3., 4.]])
print(v @ m1)
print(torch.dot(v, v))    # 向量点积:1×1 + 2×2

输出:

tensor([ 7., 10.])
tensor(5.)

v @ m1v 视为行向量 [1, 2],结果是 [1×1+2×3, 1×2+2×4] = [7, 10]torch.dot 要求两个一维向量等长,返回标量。

深度学习里 @ 随处可见:线性层就是「输入 @ 权重 + 偏置」(第 7 章)。

4.5 广播

广播(broadcasting)是形状不完全相同、但满足一定规则的张量也能做逐元素运算的机制。规则如下:

  1. 最后一个维度往前对齐;
  2. 每个维度上,两个张量要么长度相等,要么其中一个长度是 1(或缺失,视为 1);
  3. 长度为 1 的维度会被「复制扩展」到另一个的长度。

最常见的场景是「矩阵 + 行向量」:

python
import torch

base = torch.tensor([[1., 2., 3.],
                     [4., 5., 6.]])      # 形状 (2, 3)
bonus = torch.tensor([10., 20., 30.])    # 形状 (3,)
print(base + bonus)

输出:

tensor([[11., 22., 33.],
        [14., 25., 36.]])

对齐过程:bonus 形状 (3,) 与 base 形状 (2, 3) 从右往左看,第 1 维都是 3,匹配;第 0 维 base 是 2,bonus「缺失」视为 1,于是 bonus 被扩展成 2 行,每行都是 [10, 20, 30]

标量与张量运算也是广播:base * 2 中,标量 2 被扩展到整个形状:

python
print(base * 2)

输出:

tensor([[ 2.,  4.,  6.],
        [ 8., 10., 12.]])

4.6 广播失败:形状不匹配

当两个维度的长度既不相等、也没有 1 时,广播失败:

>>> import torch
>>> torch.tensor([[1., 2.], [3., 4.]]) + torch.tensor([1., 2., 3.])
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
RuntimeError: The size of tensor a (2) must match the size of tensor b (3) at non-singleton dimension 1

解读:第 1 维(从 0 数)上,a 的大小是 2,b 的大小是 3,两个都不为 1,无法广播。读报错的习惯:先看维度编号,再看该维度上两个张量各是多少

判断能不能广播的口诀:形状从右往左对齐,每个维度「相等或有一方为 1」。

动手实践

  1. 创建两个长度 5 的整数张量,分别做 +*//% 并打印结果。
  2. torch.tensor([[1., 2.], [3., 4.]]) 分别做 * 3@ torch.tensor([[1., 0.], [0., 1.]]),解释两个结果的差异。
  3. 用广播给一个 3×4 的「月销售额」矩阵每列加上该列的「基准线」向量(长度 4),输出结果。
  4. 预测 torch.tensor([[1., 2.], [3., 4.]]) + torch.tensor([[10.], [20.]]) 的结果,再运行验证,并用广播规则解释。
  5. 故意构造一个广播失败的例子,用「读报错」步骤说出:哪个维度不匹配、两个长度各是多少。

常见错误

错误写法现象原因
(a > 1) & (a < 3) 写成 a > 1 & a < 3TypeError 或结果错误& 优先级高于比较运算符,条件要加括号
需要矩阵乘法却写 m1 * m2得到逐元素乘积* 是逐元素运算,矩阵乘法用 @
a / 2 期望整数结果得到浮点除法总是返回浮点;要整数用 //
两个不兼容形状直接相加RuntimeError: The size of tensor a ... must match ...不满足广播规则,先查 shape
对需要求梯度的张量用 x.add_(10)后续 backward() 报错 a leaf Variable that requires grad is being used in an in-place operation原地运算破坏计算图,优先用 x + 10

章末练习

基础

  1. 创建 a = torch.tensor([1, 2, 3, 4])b = torch.tensor([10, 20, 30, 40]),打印 a + bb - aa * bb / a
  2. 用比较运算找出 torch.tensor([5, 12, 7, 18, 3]) 中所有大于 10 的元素(输出布尔张量,再用布尔索引取出)。
  3. 计算 torch.tensor([[1., 2.], [3., 4.]]) @ torch.tensor([[5., 6.], [7., 8.]]),手算验证第一个元素。

提高

  1. 解释广播规则,并判断以下运算哪些能成功、结果形状是什么:(2, 3) + (3,)(2, 3) + (2,)(3, 1) + (1, 3)(2, 3) + (3, 2)。然后逐一运行验证。
  2. 给形状 (4, 1) 的「每日增幅」矩阵加一个形状 (1, 3) 的「基准价」矩阵,解释广播后结果形状,并用它模拟 4 天 3 种商品的每日价格。

挑战

  1. @ 实现「3 个学生 × 4 门课成绩」矩阵与「4 门课权重」向量的加权总分,输出每个学生的总分。
  2. 不运行代码,预测 torch.tensor([[1., 2.], [3., 4.]]) + torch.tensor([10., 20.]) 的结果,运行验证后说明这个例子是「矩阵 + 行向量」还是「矩阵 + 列向量」,为什么。

章末自测

每题选择一个最佳答案。本书不附答案:完成后交由老师或 AI 老师批改讲解。

  1. torch.tensor([1, 2, 3]) * 2 的结果是?
    • A. tensor([1, 2, 3, 1, 2, 3])
    • B. tensor([2, 4, 6])
    • C. 2
    • D. 报错
  2. torch.tensor([1, 2, 3]) / 2 的结果是?
    • A. tensor([0, 1, 1])
    • B. tensor([0.5000, 1.0000, 1.5000])
    • C. tensor([1, 2, 3])
    • D. 报错
  3. 矩阵乘法应使用哪个运算符?
    • A. *
    • B. @
    • C. &
    • D. #
  4. torch.tensor([[1., 2.], [3., 4.]]) @ torch.tensor([[5., 6.], [7., 8.]]) 的第一个元素是?
    • A. 5
    • B. 19
    • C. 23
    • D. 35
  5. 形状 (2, 3)(3,) 的张量相加,结果形状是?
    • A. (2,)
    • B. (2, 3)
    • C. (3, 2)
    • D. 报错
  6. 形状 (2, 3)(3, 2) 的张量相加,会?
    • A. 得到形状 (2, 2) 的结果
    • B. 得到形状 (2, 3) 的结果
    • C. 报 RuntimeError
    • D. 得到形状 (3, 2) 的结果
  7. torch.tensor([1, 2, 3]) > 2 的结果是?
    • A. tensor([True, True, True])
    • B. tensor([False, False, True])
    • C. True
    • D. tensor([1, 2, 3])
  8. x.add_(10)x + 10 的区别是?
    • A. 没有区别
    • B. add_ 原地修改 x,x + 10 返回新张量
    • C. x + 10 原地修改 x
    • D. add_ 只能用于整数张量
  9. torch.dot(v, v) 要求 v 是?
    • A. 任意形状张量
    • B. 两个等长的一维张量
    • C. 两个二维矩阵
    • D. 标量
  10. 下列哪种情况满足广播规则?
    • A. (2, 3)(3, 2)
    • B. (2, 3)(3,)
    • C. (2, 3)(4, 3)
    • D. (5, 1)(1, 5, 5)