主要记录下torch.matmul(A,B)的用法中的一种情况:
当A,B有一个是3维以上,另一个是3维或3维以上时,如果想要使用torch.matmul(A,B),必须同时满足:
1.A和B的最后两个维度满足矩阵乘法的要求。例如A的维度是(3,1,3,3),B是(3,3,2),此时A的最后2维是(3,3),B是(3,2),符合条件
2.除去最后两个维度,A和B的其他维度要满足可以广播的条件。例如A的维度是(3,1,3,3),B是(3,3,2),除去最后两个维度,A剩下了(3,1),B是(3,),很明显此时满足广播条件
广播机制参考:Pytorch中的广播机制(Broadcast)_pytorch广播机制-CSDN博客
满足上述要求后,继续看如何执行 torch.matmul(A,B):
还是对于A的维度是(3,1,3,3),B是(3,3,2)的这个例子,假设相乘后的结果为C,那么C的维度可以按下面步骤求出:
1.A和B的最后两个维度执行矩阵乘法,得到的维度就是C的最后两个维度,也就是(3,2)
2.除去A和B的最后两个维度,其他维度进行广播后得到的维度是(3,3),这就是C的前几个维度
3.将上述维度合起来就是C的维度:(3,3,3,2)
关于torch.matmul(A,B)的更多用法参考:
Pytorch中张量矩阵乘法函数(mm, bmm, matmul)使用说明,含高维张量实例及运行结果_torch.mm-CSDN博客