掌握 PyTorch 张量乘法：八个关键函数与应用场景对比解析

PyTorch提供了几种张量乘法的方法，每种方法都是不同的，并且有不同的应用。我们来详细介绍每个方法，并且详细解释这些函数有什么区别：

1、torch.matmul

torch.matmul 是 PyTorch 中用于矩阵乘法的函数。它能够处理各种不同维度的张量，并根据张量的维度自动调整其操作方式。

torch.matmul 可以执行以下几种矩阵乘法：

二维张量之间的矩阵乘法：
- 这是经典的矩阵乘法操作。当两个张量都是二维的 (即矩阵)，torch.matmul 进行标准的矩阵乘法操作。
- 例如：假设 A 是形状为 (m, n) 的张量，B 是形状为 (n, p) 的张量，那么 torch.matmul(A, B) 结果是一个形状为 (m, p) 的张量。
高维张量之间的矩阵乘法：
- torch.matmul 可以处理更高维的张量。当输入张量的维度大于2时，它将执行批量矩阵乘法。
- 对于形状为 (..., m, n) 的张量 A 和形状为 (..., n, p) 的张量 B，torch.matmul(A, B) 的结果是形状为 (..., m, p) 的张量，其中 ... 表示相同的批量维度。批量维度部分将自动广播。
一维和二维张量的乘法：
- 当第一个张量是1D张量（向量），第二个张量是2D张量时，torch.matmul 会将1D张量视为行向量（或列向量）参与矩阵乘法。
- 例如：A 是形状为 (n,) 的张量，B 是形状为 (n, p) 的张量，那么 torch.matmul(A, B) 的结果是形状为 (p,) 的张量。
- 反之，如果第一个张量是2D张量，第二个是1D张量，则结果是一个形状为 (m,) 的张量。

import torch

# 示例 1: 二维张量之间的矩阵乘法
A = torch.tensor([[1, 2], [3, 4]])
B = torch.tensor([[5, 6], [7, 8]])
result = torch.matmul(A, B)
print(result)  # 输出: tensor([[19, 22], [43, 50]])

# 示例 2: 高维张量之间的矩阵乘法（批次矩阵乘法）
A = torch.rand(2, 3, 4)
B = torch.rand(2, 4, 5)
result = torch.matmul(A, B)
print(result.shape)  # 输出: torch.Size([2, 3, 5])

# 示例 3: 1D 和 2D 张量之间的乘法
A = torch.tensor([1, 2, 3])
B = torch.tensor([[4, 5], [6, 7], [8, 9]])
result = torch.matmul(A, B)
print(result)  # 输出: tensor([40, 46])

torch.matmul 支持广播，这意味着当输入张量的形状不完全匹配时，它可以自动扩展维度以进行相应的矩阵乘法。例如，两个张量的形状分别为 (1, 2, 3) 和 (3, 4)，torch.matmul 可以将第二个张量自动扩展为形状 (1, 3, 4)，然后进行批次矩阵乘法。

torch.matmul 底层使用了高效的线性代数库（如 BLAS），确保了矩阵乘法的性能。对于大型矩阵运算，torch.matmul 通常是非常高效的。它的灵活性和性能使得它成为 PyTorch 中广泛使用的操作之一。

2、torch.mm

torch.mm 是 PyTorch 中专门用于二维张量（矩阵）之间进行矩阵乘法的函数。与 torch.matmul 不同，torch.mm 仅适用于2D张量，并且不支持高维张量或广播操作。

torch.mm 进行标准的矩阵乘法操作，适用于两个2D张量（矩阵）之间的乘法。对于形状为 (m, n) 的张量 A 和形状为 (n, p) 的张量 B，torch.mm(A, B) 的结果是一个形状为 (m, p) 的张量。

import torch

# 示例 1: 二维张量之间的矩阵乘法
A = torch.tensor([[1, 2, 3], [4, 5, 6]])
B = torch.tensor([[7, 8], [9, 10], [11, 12]])
result = torch.mm(A, B)
print(result)  # 输出: tensor([[ 58,  64], [139, 154]])

在这个例子中，矩阵 A 的形状是 (2, 3)，矩阵 B 的形状是 (3, 2)。结果矩阵的形状是 (2, 2)，且每个元素是通过对应行与列元素的乘积之和计算得出的。

torch.mm 不支持广播机制，这意味着两个输入矩阵的形状必须严格匹配（即第一个矩阵的列数必须等于第二个矩阵的行数）。

torch.mm 是针对二维矩阵乘法优化的，它利用了底层的高效线性代数库（如 BLAS）。当仅需要进行2D张量的矩阵乘法时，torch.mm 可能比 torch.matmul 更加高效，因为它避免了 torch.matmul 中针对高维张量所做的额外处理。

注意事项：

输入张量必须是二维的。如果输入是高维张量，使用 torch.mm 会导致错误。两个矩阵的形状必须是兼容的，即第一个矩阵的列数必须等于第二个矩阵的行数，否则会抛出维度不匹配的错误。

import torch

A = torch.tensor([[1, 2], [3, 4]])
B = torch.tensor([1, 2])
# 这会引发一个错误，因为 B 不是二维张量
result = torch.mm(A, B)  # RuntimeError: matrices expected, got 1D, 2D tensors

在上面的示例中，由于 B 是一维张量而非二维矩阵，因此 torch.mm 会抛出错误。解决方法是将 B 转换为二维张量，例如 B.unsqueeze(1)，以使其形状符合矩阵乘法的要求。

torch.mm 常用于涉及矩阵乘法的各种场景，特别是在机器学习和深度学习中。例如，在神经网络的全连接层中，计算权重矩阵和输入向量的乘积时经常使用 torch.mm。此外，torch.mm 也可以用于线性代数中的基本操作，如求解线性方程组、计算特征值等。

torch.mm 它操作简洁且性能高效，适用于需要进行标准矩阵乘法的场景。对于二维矩阵乘法来说，它比 torch.matmul 更直接，因此在需要矩阵乘法且确定张量维度为2D的情况下，torch.mm 是一个理想的选择。

3、torch.bmm

torch.bmm 是 PyTorch 中用于进行批次矩阵乘法的函数。它专门处理三维张量，其中第一个维度表示批次大小，后两个维度表示需要进行矩阵乘法的矩阵。因此torch.bmm 是进行批次矩阵操作的一个高效工具。

torch.bmm 用于对形状为 (b, m, n) 的张量 A 和形状为 (b, n, p) 的张量 B 进行批次矩阵乘法，输出结果是形状为 (b, m, p) 的张量。这里，b 表示批次大小，m 和 n 是矩阵的行和列数，p 是结果矩阵的列数。

import torch

# 示例: 批次矩阵乘法
A = torch.randn(10, 3, 4)  # 形状为 (10, 3, 4)
B = torch.randn(10, 4, 5)  # 形状为 (10, 4, 5)
result = torch.bmm(A, B)
print(result.shape)  # 输出: torch.Size([10, 3, 5])

在这个例子中：

张量 A 的形状是 (10, 3, 4)，表示有10个3x4的矩阵。
张量 B 的形状是 (10, 4, 5)，表示有10个4x5的矩阵。
torch.bmm(A, B) 的结果是形状为 (10, 3, 5) 的张量，这表示批次中的每一对矩阵都进行了乘法操作。

torch.bmm 实际上是对批次中的每一对矩阵单独进行矩阵乘法操作，因此它要求输入张量的第一个维度（即批次大小）是相同的，并且后两个维度必须满足矩阵乘法的要求（即第一个矩阵的列数等于第二个矩阵的行数）。

torch.bmm 对批次矩阵乘法进行了优化，使用了高效的底层线性代数库。它在处理大型批次矩阵乘法时性能非常高效。由于它可以在批次上并行执行操作，因此特别适用于深度学习中的批量计算场景。

torch.bmm 只适用于三维张量，其中第一个维度表示批次大小。对于高于或低于三维的张量，它会报错。或者说他是torch.mm的批次化版本。torch.bmm 不支持广播机制，因此输入张量的第一个维度（批次大小）必须严格相同。

torch.bmm 常用于需要对多个矩阵对同时进行乘法操作的场景，特别是在深度学习中的以下情境：

批量计算：在训练神经网络时，我们通常将输入数据分批处理，每批次数据对应多个矩阵。torch.bmm 可以有效地处理这种批次矩阵操作。
图卷积网络（GCN）：在图神经网络中，批次矩阵乘法经常用于计算节点特征和邻接矩阵的乘积。
时间序列模型：在时间序列建模中，可能需要对每个时间步长应用不同的变换矩阵，这时可以使用 torch.bmm 进行批量处理。

torch.bmm 是专门用于批次矩阵乘法。当需要对多个矩阵对同时进行乘法操作时，它提供了高效且简洁的解决方案。

4、torch.mul

torch.mul 是 PyTorch 中用于执行元素级乘法（也称为逐元素乘法）的函数。它可以对张量的每个元素进行对应位置的乘法操作，支持任意维度的张量，并且可以自动进行广播操作来适应不同形状的张量。

torch.mul 可以对两个张量的对应元素进行乘法运算。假设有两个张量 A 和 B，那么 torch.mul(A, B) 将返回一个新的张量，其中每个元素是 A 和 B 在相同位置的元素的乘积。这个操作等同于使用 * 操作符，如 A * B。

import torch

# 示例 1: 相同形状的张量的元素级乘法
A = torch.tensor([1, 2, 3])
B = torch.tensor([4, 5, 6])
result = torch.mul(A, B)
print(result)  # 输出: tensor([ 4, 10, 18])

# 示例 2: 不同形状的张量进行广播后的元素级乘法
A = torch.tensor([[1, 2, 3], [4, 5, 6]])
B = torch.tensor([10, 20, 30])
result = torch.mul(A, B)
print(result)  # 输出: tensor([[10, 40, 90], [40, 100, 180]])

# 示例 3: 通过标量进行元素级乘法
A = torch.tensor([1, 2, 3])
result = torch.mul(A, 10)
print(result)  # 输出: tensor([10, 20, 30])

在这些示例中：

在第一个示例中，A 和 B 是形状相同的张量，因此对应元素直接相乘。
在第二个示例中，A 是二维张量，而 B 是一维张量，PyTorch 自动对 B 进行广播，使其形状与 A 匹配，然后进行逐元素乘法。
在第三个示例中，A 和一个标量值相乘，每个元素都乘以该标量。

torch.mul 支持广播机制，这意味着当两个张量的形状不完全相同时，它可以自动扩展较小形状的张量，使其与较大形状的张量兼容，然后进行逐元素乘法。

import torch

A = torch.tensor([[1, 2, 3], [4, 5, 6]])
B = torch.tensor([10, 20, 30])
result = torch.mul(A, B)

在这个例子中，A 的形状是 (2, 3)，而 B 的形状是 (3,)。PyTorch 自动将 B 扩展为 (2, 3)，然后对每个对应元素进行乘法运算。

torch.mul 是一个高效的逐元素操作，因为它直接在元素级别上进行计算，适用于需要对大批量数据进行逐元素操作的场景。它可以充分利用现代硬件的并行计算能力（如GPU），在处理大型张量时非常高效。

注意事项

虽然 torch.mul 支持广播，但在进行操作时，确保两个张量的形状是兼容的非常重要。如果形状不兼容，将会引发运行时错误。当使用标量时，标量会被自动广播到张量的每个元素，因此直接操作是安全的。

import torch

A = torch.tensor([1, 2, 3])
B = torch.tensor([[1, 2, 3], [4, 5, 6]])
# 形状不兼容，无法进行逐元素乘法
result = torch.mul(A, B)  # 会引发 RuntimeError: The size of tensor a (3) must match the size of tensor b (2) at non-singleton dimension 0

在这个错误示例中，由于 A 是一维张量，而 B 是二维张量且第一个维度不匹配，因此无法广播，导致错误。

torch.mul 在许多机器学习和深度学习任务中都非常有用。例如：

权重调整：在神经网络中，可以通过 torch.mul 来逐元素调整权重或激活值。
掩码操作：在图像处理中，可以使用 torch.mul 来对图像应用掩码，逐元素控制哪些部分需要保留或修改。
归一化：可以逐元素将张量归一化或缩放，以满足特定的算法要求。

torch.mul 在处理各种张量操作时非常有用。它支持广播机制，可以自动适应不同形状的张量，从而在多种应用场景中提供简洁而高效的解决方案。

5、torch.mv

torch.mv 是 PyTorch 中用于进行矩阵与向量乘法的函数。它专门用于二维张量（矩阵）和一维张量（向量）之间的乘法操作。torch.mv 是矩阵乘法的一种特殊情况，适用于当你需要将矩阵乘以向量时使用。

torch.mv 执行的是矩阵与向量的乘法操作。假设有一个矩阵 A，它的形状为 (m, n)，以及一个向量 v，它的形状为 (n,)，那么 torch.mv(A, v) 将返回一个形状为 (m,) 的一维张量（向量），结果是矩阵 A 与向量 v 的乘积。

import torch

# 示例: 矩阵与向量的乘法
A = torch.tensor([[1, 2, 3], [4, 5, 6]])
v = torch.tensor([7, 8, 9])
result = torch.mv(A, v)
print(result)  # 输出: tensor([ 50, 122])

在这个示例中,矩阵 A 的形状为 (2, 3)，向量 v 的形状为 (3,)。通过 torch.mv(A, v)，我们得到的结果是形状为 (2,) 的向量 [50, 122]，其中每个元素是通过矩阵与向量的标准乘法计算得出的。

torch.mv 执行的矩阵与向量乘法遵循以下规则：对于矩阵 A 中的每一行，将该行与向量 v 的所有元素逐元素相乘，并将乘积的结果求和，得到一个标量。这个标量就是结果向量对应位置的值。

import torch

A = torch.tensor([[1, 2, 3], [4, 5, 6]])
v = torch.tensor([7, 8, 9])
result = torch.mv(A, v)
# 结果:
# result[0] = 1*7 + 2*8 + 3*9 = 50
# result[1] = 4*7 + 5*8 + 6*9 = 122

torch.mv 专门用于矩阵和向量的乘法，比通用的矩阵乘法函数如 torch.matmul 或 torch.mm 更加高效，因为它避免了对多余维度的处理。这使得 torch.mv 在执行矩阵与向量乘法时速度更快，并且更适合用于大规模计算。

注意事项

矩阵 A 的列数（第二个维度）必须等于向量 v 的长度（第一个维度），否则将会报错。

import torch

A = torch.tensor([[1, 2, 3], [4, 5, 6]])
v = torch.tensor([7, 8])
# 这将引发错误，因为 v 的形状与 A 的列数不匹配
result = torch.mv(A, v)  # 会引发 RuntimeError: size mismatch, m1: [2x3], m2: [2] at THTensorMath.cpp:41

在这个错误示例中，向量 v 的长度与矩阵 A 的列数不匹配，因此无法进行矩阵与向量乘法。

torch.mv 是 PyTorch 中用于执行矩阵与向量乘法的专用函数。它对矩阵与向量乘法进行了优化，能够高效处理这类操作，是线性代数、深度学习和科学计算中常用的工具。在许多应用场景中都很有用，特别是在以下情况下：

线性代数操作：在计算线性方程组、特征值问题等线性代数问题时，经常需要进行矩阵与向量的乘法。
神经网络计算：在神经网络的前向传播过程中，特别是全连接层中，权重矩阵与输入向量的乘法操作可以通过 torch.mv 高效地实现。
物理模拟：在一些物理模拟中，状态向量与转换矩阵的乘法操作可以通过 torch.mv 实现。

6、torch.dot

torch.dot 是 PyTorch 中用于计算两个一维张量（即向量）之间的点乘（内积）的函数。点乘是一种基本的向量操作，在许多数学和工程应用中都有广泛的应用。

torch.dot 计算的是两个向量之间的点积。假设有两个向量 a 和 b，它们的长度相同（即形状都为 (n,)），那么 torch.dot(a, b) 的结果是一个标量（即一个数值），这个值是通过对应位置的元素相乘后再求和得到的。

import torch

# 示例: 两个向量的点乘
a = torch.tensor([1, 2, 3])
b = torch.tensor([4, 5, 6])
result = torch.dot(a, b)
print(result)  # 输出: tensor(32)

在这个示例中：向量 a 的形状为 (3,)，向量 b 的形状也是 (3,)。通过 torch.dot(a, b)，我们得到了标量 32，其计算过程为：1*4 + 2*5 + 3*6 = 4 + 10 + 18 = 32。

torch.dot 计算点乘的方式是逐元素相乘，然后将结果求和。对于两个长度为 n 的向量 a 和 b，点积的计算公式如下：

result = (a[0] * b[0]) + (a[1] * b[1]) + ... + (a[n-1] * b[n-1])

torch.dot 是对两个一维张量进行点积的优化实现，由于其简单的计算流程和对向量操作的专门优化，它通常具有非常高的性能，特别是在 GPU 上处理大规模数据时表现尤为优异。

torch.dot 仅适用于一维张量（向量），如果输入的张量不是一维的，会引发错误。并且torch.dot 返回一个标量（标量张量），而不是张量。由于点积的对称性，torch.dot(a, b) 与 torch.dot(b, a) 的结果是相同的。

与其他操作的对比

torch.matmul 和 torch.mm：这些函数用于矩阵乘法，适用于高维张量。torch.dot 只用于一维张量的点积。
torch.mul：这是逐元素乘法，不是点积。torch.mul(a, b) 会返回一个与 a 和 b 形状相同的张量，其中每个元素是对应元素的乘积，而 torch.dot(a, b) 会返回一个标量。

torch.dot 是一个简单而高效的函数，专门用于计算一维张量之间的点积。在许多数学、物理和工程应用中，它是一个非常重要的工具。点积在很多场景中都有应用，包括但不限于：

向量投影：在几何中，点乘可以用于计算一个向量在另一个向量方向上的投影。
相似性计算：在信息检索和机器学习中，两个向量的点积可以用于衡量它们的相似性。例如，在词向量（Word Embeddings）的相似性计算中，点积是常用的度量方法之一。
能量计算：在物理学中，点积用于计算力和位移的乘积（即功的计算）。

7、torch.outer

torch.outer 是 PyTorch 中用于计算两个一维张量（即向量）之间的外积（外积矩阵）的函数。外积是线性代数中的一种基本运算，结果是一个矩阵，其元素是两个输入向量各元素的乘积。

torch.outer 计算的是两个向量的外积。假设有两个向量 a 和 b，它们的形状分别是 (n,) 和 (m,)，那么 torch.outer(a, b) 的结果是一个形状为 (n, m) 的二维张量（矩阵），这个矩阵中的元素由 a[i] * b[j] 计算得到。

import torch

# 示例: 两个向量的外积
a = torch.tensor([1, 2, 3])
b = torch.tensor([4, 5, 6])
result = torch.outer(a, b)
print(result)
# 输出:
# tensor([[ 4,  5,  6],
#         [ 8, 10, 12],
#         [12, 15, 18]])

在这个示例中：

向量 a 的形状为 (3,)，向量 b 的形状也为 (3,)。
通过 torch.outer(a, b)，我们得到了形状为 (3, 3) 的矩阵。这个矩阵的每个元素都是由 a[i] 和 b[j] 的乘积计算得出。

torch.outer 是对两个一维张量进行外积的优化实现。由于其操作涉及大量的元素乘法，因此在处理大型向量时，特别是在 GPU 上计算，torch.outer 的性能表现十分出色。

torch.outer 仅适用于一维张量，即向量，并返回一个二维张量（矩阵），其形状为 (n, m)，其中 n 和 m 是输入向量的长度。

与其他操作的对比

torch.matmul 和 torch.mm：这些函数用于矩阵乘法，适用于高维张量。torch.outer 专用于计算两个一维张量之间的外积。
torch.mul：这是逐元素乘法。如果两个张量的形状相同，torch.mul(a, b) 将执行逐元素乘法，而不是计算外积。

torch.outer 是一个用于计算两个一维张量之间外积的高效工具。它在生成矩阵、处理双线性形式、构建张量积等应用中非常有用。外积在很多场景中都有应用，包括但不限于：

矩阵构建：外积可用于生成特定类型的矩阵，例如克罗内克积。
双线性形式：在双线性形式的表示中，外积经常用于构建张量。
机器学习：在神经网络的权重更新、特征交互等场景中，外积运算可以构造高阶特征。

8、torch.einsum

torch.einsum 是 PyTorch 中一个非常强大的函数，它使用爱因斯坦求和约定（Einstein Summation Convention）来执行复杂的张量操作。torch.einsum 的灵活性使得它可以用于各种矩阵和张量运算，包括矩阵乘法、转置、内积、外积、以及其他高阶张量运算。

爱因斯坦求和约定是一种简化张量操作的符号表示方法，其中重复的指标自动表示求和。torch.einsum 使用字符串表示张量操作，将输入张量的维度与输出维度通过指定的模式进行映射。

torch.einsum(equation, *operands)

equation：一个字符串，描述了输入和输出张量的维度关系。

*operands：一个或多个张量，参与计算的张量。

使用示例

1、矩阵乘法

矩阵乘法是最常见的张量操作之一。对于两个矩阵 A 和 B，使用 torch.einsum 进行矩阵乘法可以表示为：

import torch

A = torch.tensor([[1, 2], [3, 4]])
B = torch.tensor([[5, 6], [7, 8]])
result = torch.einsum('ik,kj->ij', A, B)
print(result)  # 输出: tensor([[19, 22], [43, 50]])

这里，'ik,kj->ij' 表示：

A 的维度为 i（行）和 k（列）。
B 的维度为 k（行）和 j（列）。
输出的矩阵 C 的维度为 i（行）和 j（列），其中 k 是求和维度。

2、向量内积（点积）

对于两个向量 a 和 b，它们的内积可以用 torch.einsum 表示为：

a = torch.tensor([1, 2, 3])
b = torch.tensor([4, 5, 6])
result = torch.einsum('i,i->', a, b)
print(result)  # 输出: tensor(32)

这里，'i,i->' 表示：

a 和 b 都是一维向量，维度为 i。
输出是一个标量（没有索引），表示所有元素的乘积之和。

3、向量外积

向量外积可以表示为：

a = torch.tensor([1, 2, 3])
b = torch.tensor([4, 5, 6])
result = torch.einsum('i,j->ij', a, b)
print(result)
# 输出:
# tensor([[ 4,  5,  6],
#         [ 8, 10, 12],
#         [12, 15, 18]])

这里，'i,j->ij' 表示：

a 的维度为 i，b 的维度为 j。
输出矩阵 C 的维度为 ij，表示 a[i] 和 b[j] 的乘积。

torch.einsum 是一个通用且灵活的工具，但其性能可能不如专门为某些操作优化的函数（如 torch.matmul）。所以在性能关键的应用中，使用专门的张量操作函数可能会更高效。不过对于需要简洁表示复杂操作的场景，torch.einsum 仍然是首选。

总结

以下是对 PyTorch 中几种常用张量操作函数的总结：

torch.matmul (矩阵乘法)

功能：执行矩阵乘法，支持二维矩阵、批量矩阵乘法、高维张量乘法。
应用：广泛用于神经网络中的矩阵运算，如全连接层的计算。

torch.mm (矩阵乘法)

功能：专门用于二维张量（矩阵）之间的乘法，不支持广播和高维张量。
应用：适用于明确为二维矩阵的乘法操作，性能高效。

torch.bmm (批次矩阵乘法)

功能：对三维张量进行批次矩阵乘法，适用于批量处理的场景。
应用：常用于深度学习中的批量数据处理和图神经网络中的邻接矩阵计算。

torch.mul (元素级乘法)

功能：逐元素乘法，支持任意维度张量并自动广播。
应用：用于权重调整、掩码操作、数据归一化等逐元素运算。

torch.mv (矩阵与向量乘法)

功能：用于二维矩阵与一维向量之间的乘法操作。
应用：适用于神经网络中的前向传播、线性代数操作。

torch.dot (点乘)

功能：计算两个一维张量（向量）之间的点积，结果是一个标量。
应用：用于计算向量内积、向量相似性、物理学中的能量计算。

torch.outer (外积)

功能：计算两个一维张量之间的外积，结果是一个二维矩阵。
应用：用于构建矩阵、处理双线性形式、特征交互等。

torch.einsum (爱因斯坦求和约定)

功能：使用爱因斯坦求和约定进行复杂张量运算，包括矩阵乘法、转置、内积、外积等。
应用：广泛用于线性代数、物理学计算、机器学习中的复杂操作。

这些 PyTorch 张量操作函数各有其专门用途和应用场景。torch.matmul、torch.mm 和 torch.bmm 主要用于矩阵乘法；torch.mul 和 torch.outer 用于逐元素和外积操作；torch.mv 和 torch.dot 处理矩阵与向量、向量与向量的乘法；torch.einsum 则是处理复杂张量运算的多功能工具。

https://avoid.overfit.cn/post/2ce5023db7634886b602e7eb6ce4abb7