torch.bmm
是 PyTorch
中的一个函数,用于执行批量矩阵相乘(batch matrix multiplication
)的操作。它用于计算两个具有相同批次大小的三维张量的矩阵乘法。
在矩阵乘法中,两个矩阵的维度必须满足一定的条件。对于 torch.bmm
函数,它要求输入的两个张量都具有三个维度,形状分别为 (batch_size, n, m)
和 (batch_size, m, p)
,其中 batch_size
表示批次大小,n、m 和 p
分别表示矩阵的行数和列数。
torch.bmm
函数将执行批量矩阵相乘的操作,计算每个批次中对应位置的两个矩阵的乘积。它返回一个具有形状 (batch_size, n, p)
的新张量,其中每个批次中的结果矩阵是对应位置的输入矩阵相乘的结果。
以下是 torch.bmm 的使用示例:
import torch
# 创建两个具有相同批次大小的三维张量
batch_size = 2
n = 3
m = 4
p = 5
x = torch.randn(batch_size, n, m)
y = torch.randn(batch_size, m, p)
# 执行批量矩阵相乘操作
result = torch.bmm(x, y)
# 打印结果张量的形状
print(result.shape)
在这个示例中,我们创建了两个具有相同批次大小的三维张量 x
和 y
。它们的形状分别为 (2, 3, 4)
和 (2, 4, 5)
。然后,我们使用 torch.bmm
对这两个张量进行批量矩阵相乘操作,并将结果存储在 result
张量中。最后,打印出结果张量的形状是 (2, 3, 5)
。
torch.bmm
要求输入张量的维度满足特定的条件,并且批次大小必须相同。如果输入的张量不满足要求,将会引发错误。因此,在使用 torch.bmm
之前,请确保输入张量的维度和批次大小是符合要求的。