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])