torch.bmmPyTorchis a function in that performs batch matrix multiplication ( batch matrix multiplication). It is used to compute the matrix multiplication of two 3D tensors with the same batch size.
In matrix multiplication, the dimensions of the two matrices must satisfy certain conditions. For torch.bmmthe function , it requires that the two input tensors have three dimensions and are of shape (batch_size, n, m)and respectively (batch_size, m, p), where batch_sizerepresents the batch size and n、m 和 prepresents the number of rows and columns of the matrix, respectively.
torch.bmmThe function will perform batch matrix multiplication, calculating the product of two matrices at corresponding positions in each batch. It returns a new tensor (batch_size, n, p)with where the resulting matrix in each batch is the result of multiplying the input matrices at the corresponding positions.
Here is an example of using 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)
In this example, we create two 3D tensors with the same batch size xand y. Their shapes are (2, 3, 4)and , (2, 4, 5). We then perform a batch matrix multiplication operation torch.bmmon these and store the result resultin the tensor. Finally, print out what the shape of the resulting tensor is (2, 3, 5).
torch.bmmRequires that the dimensions of the input tensors meet certain conditions and that the batch sizes must be the same. If the input tensor does not meet the requirements, an error will be raised. Therefore, torch.bmmbefore , please make sure that the dimensions and batch size of the input tensors meet the requirements.