Python PyTorch bmm用法及代码示例

本文简要介绍python语言中 torch.bmm 的用法。

用法:

torch.bmm(input, mat2, *, out=None) → Tensor

参数：

• input(Tensor) -第一批要相乘的矩阵

• mat2(Tensor) -第二批要相乘的矩阵

关键字参数：

out(Tensor,可选的) -输出张量。

执行存储在 input 和 mat2 中的矩阵的批处理 matrix-matrix 乘积。

input 和 mat2 必须是 3-D 张量，每个张量都包含相同数量的矩阵。

如果input是(b \times n \times m)张量，mat2是(b \times m \times p)张量，out将是(b \times n \times p)张量。

\text{out}_i = \text{input}_i \mathbin{@} \text{mat2}_i

该运算符支持 TensorFloat32。

注意

此函数不广播。对于广播矩阵产品，请参阅 torch.matmul() 。

例子：

>>> input = torch.randn(10, 3, 4)
>>> mat2 = torch.randn(10, 4, 5)
>>> res = torch.bmm(input, mat2)
>>> res.size()
torch.Size([10, 3, 5])