Pytorch常用乘法函数总结：torch.mul()、*、torch.mm()、torch.bmm()、torch.mv()、torch.dot()、@、torch.matmul()

上一篇博客总结了numpy中常用的乘法函数：

numpy常用乘法函数总结：np.dot()、np.multiply()、*、np.matmul()、@、np.prod()、np.outer()-CSDN博客

主要是 np.dot()、np.multiply()、*、np.matmul()、@ 五种，其中 np.matmul() 和 @ 完全等价，np.multiply() 和 * 在输入数据类型为 np.array 时也完全等价

本文总结Pytorch常用的乘法函数，TensorFlow常用乘法函数放在下一篇文章里

目录

torch.mul() 和 *【等价，element-wise乘，可广播】

广播操作

torch.mm() 或 torch.bmm() 【矩阵乘法，前二维后三维，均不可广播】

torch.mv()【矩阵-向量乘法，不可广播】

torch.dot()【仅支持两个一维向量点积】

@【等价于 torch.dot() + torch.mv() + torch.mm()】

torch.matmul() 【矩阵乘法，可高维，可广播】

参考

---

torch.mul() 和 *【等价，element-wise乘，可广播】

torch.mul(x, y) 等价于 x*y 

• 矩阵 * 标量：矩阵中的每个元素都乘以标量
• 矩阵 * 行向量：要求矩阵的列数 = 行向量的列数
• 矩阵 * 列向量：要求矩阵的行数 = 列向量的行数
• 矩阵 * 矩阵：要求两个矩阵的维度完全相同，不同的话看是否能广播成相同shape

vec1 = torch.arange(4)
vec2 = torch.tensor([4,3,2,1])
mat1 = torch.arange(12).reshape(4,3)
mat2 = torch.arange(12).reshape(3,4)

print(vec1 * vec2)
print(mat2 * vec1)
print(mat1 * mat1)

Output:
tensor([0, 3, 4, 3])
tensor([[ 0,  1,  4,  9],
        [ 0,  5, 12, 21],
        [ 0,  9, 20, 33]])
tensor([[  0,   1,   4],
        [  9,  16,  25],
        [ 36,  49,  64],
        [ 81, 100, 121]])

一句话总结：如果两者shape相同，直接element-wise相乘；如果两者shape不同，本质上是先把另一个扩展/复制成相同的shape（即广播操作），再对应元素相乘 

广播操作

（1）在一定的规则下允许高维Tensor和低维Tensor之间的运算。这里举一个例子：a是二维Tensor，b是三维Tensor，但是a的维度与b的后两位相同，那么a和b仍然可以做 * 操作，结果是一个和b维度一样的三维Tensor。可以理解为沿着b的第0维做二维Tensor点积，或运算前将a沿着b的第0维进行了expand操作

import torch

a = torch.tensor([[1, 2], [2, 3]])   # torch.Size([2, 2])
b = torch.tensor([[[1, 2], [2, 3]], [[-1, -2], [-2, -3]]])   # torch.Size([2, 2, 2])

print(a*b)    # torch.Size([2, 2, 2])   b*a 结果相同
'''
tensor([[[ 1,  4],
         [ 4,  9]],

        [[-1, -4],
         [-4, -9]]])
'''

运算过程可以理解为：

a = a.expand(b.size())
'''
tensor([[[1, 2],
         [2, 3]],

        [[1, 2],
         [2, 3]]])
'''

a*b

（2）再举个例子：两个Tensor都是3维，但shape不一致，这种情况看是否能够广播成一致的shape

import torch

node = torch.tensor([[1, 1, 1, 0, 0], [1, 1, 1, 1, 1]])   # shape=(2,5)

node_0 = node.unsqueeze(-1)   # torch.Size([2, 5, 1])
'''
tensor([[[1],
         [1],
         [1],
         [0],
         [0]],

        [[1],
         [1],
         [1],
         [1],
         [1]]])
'''

node_1 = node.unsqueeze(1)    # torch.Size([2, 1, 5])
'''
tensor([[[1, 1, 1, 0, 0]],

        [[1, 1, 1, 1, 1]]])
'''

# 把node_0的shape 复制成node_1的shape，node_1同理，直到两者shape一致后再对位相乘
print((node_0*node_1).shape)  # torch.Size([2, 5, 5])
'''
tensor([[[1, 1, 1, 0, 0],
         [1, 1, 1, 0, 0],
         [1, 1, 1, 0, 0],
         [0, 0, 0, 0, 0],
         [0, 0, 0, 0, 0]],

        [[1, 1, 1, 1, 1],
         [1, 1, 1, 1, 1],
         [1, 1, 1, 1, 1],
         [1, 1, 1, 1, 1],
         [1, 1, 1, 1, 1]]])
'''

---

torch.mm() 或 torch.bmm() 【矩阵乘法，前二维后三维，均不可广播】

torch.mm() 用于两个二维向量之间的矩阵乘法。如果 input1 是一个n×m张量，input2 是一个 m×p张量，将会输出一个 n×p 张量

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

而 torch.bmm() 是用于两个三维张量之间的批次矩阵乘法（其中第0维为批次大小）。由于神经网络训练一般采用mini-batch，经常输入的是三维带batch矩阵。如果 input1 是一个 b×n×m 张量，input2 是一个 b×m×p 张量，将会输出一个 b×n×p 张量

>>> a = torch.ones(3,4,5)
>>> b = torch.ones(3,5,6)
>>> torch.bmm(a, b)    # torch.Size([3, 4, 6])
tensor([[[5., 5., 5., 5., 5., 5.],
         [5., 5., 5., 5., 5., 5.],
         [5., 5., 5., 5., 5., 5.],
         [5., 5., 5., 5., 5., 5.]],
 
        [[5., 5., 5., 5., 5., 5.],
         [5., 5., 5., 5., 5., 5.],
         [5., 5., 5., 5., 5., 5.],
         [5., 5., 5., 5., 5., 5.]],
 
        [[5., 5., 5., 5., 5., 5.],
         [5., 5., 5., 5., 5., 5.],
         [5., 5., 5., 5., 5., 5.],
         [5., 5., 5., 5., 5., 5.]]])

另外，这两个函数都无法广播

---

torch.mv()【矩阵-向量乘法，不可广播】

torch.mv(input, vec) 只支持矩阵和向量之间的乘法。如果 input 为 n×m，vec向量的长度为m，那么输出为 n×1的向量。不支持广播机制

In[1]: vec = torch.arange(4)
In[2]: mat = torch.arange(12).reshape(3,4)
In[3]: torch.mv(mat, vec)
Out[1]: tensor([14, 38, 62])

---

torch.dot()【仅支持两个一维向量点积】

只能支持两个一维向量，两向量相乘相加得到一个标量

• 与 np.dot() 不同，np.dot() 既能向量点积，又能矩阵乘法

a = torch.tensor([2,3])
b = torch.tensor([1,2])
c = torch.dot(a,b)
print('a:',a.shape)                 # a: torch.Size([2])
print('b:',b.shape)                 # b: torch.Size([2])
print('torch.dot:',c,c.shape)       # torch.dot: tensor(8) torch.Size([])

---

@【等价于 torch.dot() + torch.mv() + torch.mm()】

mat1 @ mat2

• 若mat1和mat2都是一维向量，那么对应操作就是torch.dot()
• 若mat1是二维向量，mat2是一维向量，那么对应操作就是torch.mv()
• 若mat1和mat2都是二维向量，那么对应操作就是torch.mm()

vec1 = torch.arange(4)
vec2 = torch.tensor([4,3,2,1])
mat1 = torch.arange(12).reshape(4,3)
mat2 = torch.arange(12).reshape(3,4)

print(vec1 @ vec2) # 两个一维向量
print(mat2 @ vec1) # 一个二维和一个一维
print(mat1 @ mat2) # 两个二维向量

Output:
tensor(10)
tensor([14, 38, 62])
tensor([[ 20,  23,  26,  29],
        [ 56,  68,  80,  92],
        [ 92, 113, 134, 155],
        [128, 158, 188, 218]])

---

torch.matmul() 【矩阵乘法，可高维，可广播】

torch.matmul() 与 @ 类似，但它不止局限于一维和二维，可以扩展到高维，且可以广播

• 如果输入的两个张量都是一维的，那么返回点积，得到一个数，对应的操作就是torch.dot()
• 如果输入的两个张量都是二维的，那么返回矩阵乘积，对应的操作就是torch.mm()
• 如果输入的第一个张量是二维的，第二个张量是一维的，matrix第二维需要与 vector维度相同，返回矩阵向量乘积，对应的操作就是torch.mv()

以上三种情况就不举例子了，见对应函数的例子即可

• 如果输入的第一个张量是一维的，第二个参数是二维的，那么torch.matmul()操作会先将第一个张量的维度前面添加1，在执行矩阵相乘后，再将添加的维度移除

import torch

'''
torch.matmul()执行过程：
    torch.Size([4]) 扩展为 torch.Size([1,4]) ->
    torch.Size([1,4]) 与 torch.Size([4,3]) 矩阵相乘 ->
    torch.Size([1,3]) 再将添加的维度移除 ->
    torch.Size([3]) 最终结果
'''
vec = torch.tensor([1,2,3,4])         # torch.Size([4])
mat = torch.arange(12).reshape(4,3)   # torch.Size([4,3])
print(torch.matmul(vec, mat))         # tensor([60, 70, 80])
print(torch.matmul(vec, mat).shape)   # torch.Size([3])

• 如果两个参数一个是一维的，另一个是N维的（N>2），则进行批处理矩阵乘法
  • 如果第一个参数一维，则在第一个参数的维度前面加1，在进行批处理矩阵相乘后再删除
  • 如果第二个参数一维，则在第二个参数的维度后面加1，在进行批处理矩阵相乘后再删除

tensor1 = torch.randn(2)
tensor2 = torch.randn(10,2,3)
print(torch.matmul(tensor1, tensor2).size())

Output:
torch.Size([10, 3])

先将 tensor2前面维度10作为batch提出来，使其变成二维（2 * 3），将 tensor1的维度前面加1，那么就变成了1*2（二维），然后维度1*2和维度2*3做矩阵乘法，得到1*3，再将tensor1添加的维度1删除，最终得到维度10*3

• 如果两个参数都是多维，永远是最后两个维度相乘，再将前面的维度补上

如第一个参数二维，第二个参数>2维：

import torch

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

同样的a和b，使用 torch.mm() 或 torch.bmm() 均会报错

import torch

a = torch.ones(3,4)
b = torch.ones(5,4,2)

print(torch.mm(a, b))    # RuntimeError: mat2 must be a matrix
print(torch.bmm(a, b))   # RuntimeError: batch1 must be a 3D tensor

又如两个参数都>2维：

tensor1 = torch.randn(5,1,5,3)
tensor2 = torch.randn(2,3,4)
print(torch.matmul(tensor1, tensor2).shape)

Output:
torch.Size([5, 2, 5, 4])

先将tensor1中多余的一维提取出来，剩下三维，将tensor1中做广播机制，变成2 * 5 * 3，接着将tensor1和tensor2中最后两维做矩阵乘法，得到5 * 4, 最终得到维度5 * 2 * 5 * 4

---

参考

torch.Tensor的乘法汇总-CSDN博客

torch中乘法整理,*&torch.mul()&torch.mv()&torch.mm()&torch.dot()&@&torch.mutmal()-CSDN博客

pytorch | pytorch常用的乘法运算_pytorch里矩阵对应元素相乘-CSDN博客