pytorch中的matmul與mm,bmm區(qū)別說明

pytorch中matmul和mm和bmm區(qū)別 matmulmmbmm結(jié)論

先看下官網(wǎng)上對(duì)這三個(gè)函數(shù)的介紹。

matmul

mm

bmm

顧名思義, 就是兩個(gè)batch矩陣乘法.

結(jié)論

從官方文檔可以看出

1、mm只能進(jìn)行矩陣乘法,也就是輸入的兩個(gè)tensor維度只能是( n × m ) (n\times m)(n×m)和( m × p ) (m\times p)(m×p)

2、bmm是兩個(gè)三維張量相乘, 兩個(gè)輸入tensor維度是( b × n × m ) (b\times n\times m)(b×n×m)和( b × m × p ) (b\times m\times p)(b×m×p), 第一維b代表batch size，輸出為( b × n × p ) (b\times n \times p)(b×n×p)

3、matmul可以進(jìn)行張量乘法, 輸入可以是高維.

補(bǔ)充：torch中的幾種乘法。torch.mm, torch.mul, torch.matmul

一、點(diǎn)乘

點(diǎn)乘都是broadcast的，可以用torch.mul(a, b)實(shí)現(xiàn)，也可以直接用*實(shí)現(xiàn)。

>>> a = torch.ones(3,4)
>>> a
tensor([[1., 1., 1., 1.],
        [1., 1., 1., 1.],
        [1., 1., 1., 1.]])
>>> b = torch.Tensor([1,2,3]).reshape((3,1))
>>> b
tensor([[1.],
        [2.],
        [3.]])
>>> torch.mul(a, b)
tensor([[1., 1., 1., 1.],
        [2., 2., 2., 2.],
        [3., 3., 3., 3.]])

當(dāng)a, b維度不一致時(shí)，會(huì)自動(dòng)填充到相同維度相點(diǎn)乘。

二、矩陣乘

矩陣相乘有torch.mm和torch.matmul兩個(gè)函數(shù)。其中前一個(gè)是針對(duì)二維矩陣，后一個(gè)是高維。當(dāng)torch.mm用于大于二維時(shí)將報(bào)錯(cuò)。

>>> a = torch.ones(3,4)
>>> b = torch.ones(4,2)
>>> torch.mm(a, b)
tensor([[4., 4.],
        [4., 4.],
        [4., 4.]])

>>> a = torch.ones(3,4)
>>> b = torch.ones(5,4,2)
>>> torch.matmul(a, b).shape
torch.Size([5, 3, 2])

>>> a = torch.ones(5,4,2)
>>> b = torch.ones(5,2,3)
>>> torch.matmul(a, b).shape
torch.Size([5, 4, 3])

>>> a = torch.ones(5,4,2)
>>> b = torch.ones(5,2,3)
>>> torch.matmul(b, a).shape
報(bào)錯(cuò)。

以上為個(gè)人經(jīng)驗(yàn)，希望能給大家一個(gè)參考，也希望大家多多支持腳本之家。如有錯(cuò)誤或未考慮完全的地方，望不吝賜教。