PyTorch中的矩阵乘法

1. 二维矩阵乘法 

 , 其中 , , 输出 的维度是。该函数一般只用来计算两个二维矩阵的矩阵乘法，而且不支持broadcast操作。

2. 三维带Batch矩阵乘法 

由于神经网络训练一般采用mini-batch，经常输入的是三维带batch矩阵，所以提供 ,其中 , , 输出 的维度是 。该函数的两个输入必须是三维矩阵且第一维相同（表示Batch维度），不支持broadcast操作。

3. "混合"矩阵乘法 

 支持broadcast操作，使用起来比较复杂，建议参考pytorch官方文档。

特别 ，针对多维数据 乘法，我们可以认为该 乘法使用使用两个参数的后两个维度来计算，其他的维度都可以认为是batch维度。假设两个输入的维度分别是, ，那么我们可以认为  乘法首先是进行后两位矩阵乘法得到 ，然后分析两个参数的batch size分别是  和  , 可以广播成为 ， 因此最终输出的维度是 。

4. 矩阵逐元素(Element-wise)乘法 

，其中  乘数可以是标量也可以是任意维度的矩阵，只要满足最终相乘是可以broadcast的即可,即该操作是支持broadcast操作的。

•  是标量: 例如是维度任意的矩阵, (一个标量), 那么输出一个矩阵，其中每个值是 中原值乘以 , 维度保持不变。

 是矩阵: 只要  与  的维度可以满足broadcast条件，就可以进行逐元素乘法操作，例如：

1 import torch
2 A = torch.randn(2,3,4)
3 B = torch.randn(3, 4)
4 print (torch.mul(A,b).shape) # 输出 torch.size([2,3,4)

5. 两个乘法操作符@和

简单来说， @ 操作符可以执行矩阵乘法操作，类似  ; 而  乘法操作可以执行逐元素矩阵乘法，使用方法类似 。

 1 import torch
 2
 3 x=torch.ones(3,2)
 4 print(x)
 5
 6 y=torch.ones(3,2)+2
 7 print(y)
 8
 9 z=torch.ones(2,1)
10 print(z)
11
12 print(x*y@z)

参考：随笔1: PyTorch中矩阵乘法总结 - 知乎 (zhihu.com)