pytorch中matmul和mm和bmm区别


先看下官网上对这三个函数的介绍。

matmul

mm

bmm

顾名思义, 就是两个batch矩阵乘法.

结论

从官方文档可以看出,

  • mm只能进行矩阵乘法,也就是输入的两个tensor维度只能是 ( n × m ) (n\times m) (n×m) ( m × p ) (m\times p) (m×p)
  • bmm是两个三维张量相乘, 两个输入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)
  • matmul可以进行张量乘法, 输入可以是高维.