理解torch函数bmm

基本信息

功能描述

torch.bmm 是 PyTorch 中的一个函数，用于执行批量矩阵乘法（Batch Matrix Multiplication）。它适用于处理一批矩阵的乘法操作，特别适合于深度学习任务中的场景，比如卷积神经网络中的某些层。

参数说明

• input1: 第一个输入张量，形状为 (batch_size, N, M)。
• input2: 第二个输入张量，形状为 (batch_size, M, P)。
• 返回值: 一个新的张量，形状为 (batch_size, N, P)，表示每个批次内的矩阵乘法结果。

使用示例

示例1: 基本用法

import torch# 定义两个三维张量
a = torch.randn(10, 3, 4)  # 10 个 3x4 的矩阵
b = torch.randn(10, 4, 5)  # 10 个 4x5 的矩阵# 进行批次矩阵乘法
result = torch.bmm(a, b)# 输出结果并记录信息
print(f"Shape of a: {a.shape}")
print(f"Shape of b: {b.shape}")
print(f"Shape of result: {result.shape}")"""
Output:
Shape of a: torch.Size([10, 3, 4])
Shape of b: torch.Size([10, 4, 5])
Shape of result: torch.Size([10, 3, 5])
"""

在这个例子中，我们有两个形状分别为 (10, 3, 4) 和 (10, 4, 5) 的张量 a 和 b。通过调用 torch.bmm(a, b)，我们获得了一个新的张量 result，其形状为 (10, 3, 5)，这意味着对于每一个批次（共 10 个），我们都成功地完成了对应的矩阵乘法操作。

示例2: 处理单个矩阵的情况
虽然 torch.bmm 主要设计用来处理批量化矩阵乘法，但如果只有一个矩阵的话，可以通过增加额外的维度来适应这个接口。

# 单个矩阵的例子
a_single = torch.randn(3, 4).unsqueeze(0)  # 添加一个批次维度，变为 (1, 3, 4)
b_single = torch.randn(4, 5).unsqueeze(0)  # 同样添加一个批次维度，变为 (1, 4, 5)result_single = torch.bmm(a_single, b_single)print(result_single.squeeze())
"""
输出可能是一个 3x5 的矩阵，具体内容取决于随机生成的数据。
"""

在这里，我们首先将原本是二维的矩阵转换为带有单一批次维度的形式 (unsqueeze)，然后就可以直接使用 torch.bmm 来完成乘法运算。最后，如果我们只需要得到实际的结果而不关心批次维度的存在与否，可以使用 squeeze() 方法去除多余的维度。

与其他矩阵乘法函数的区别

• torch.mm vs torch.bmm:
  • torch.mm 仅支持两个二维矩阵之间的乘法。
  • torch.bmm 支持三个维度的张量，第一个维度代表批次数量，其余两个维度遵循标准的矩阵乘法规则。
• torch.matmul vs torch.bmm:
  • torch.matmul 提供更广泛的通用性，不仅限于矩阵乘法，还支持点积和其他类型的线性代数运算，并且具备广播机制。
  • torch.bmm 更专注于高效的批量矩阵乘法实现，没有广播能力，但在特定情况下性能更好。