torch.Tensor的4种乘法

torch.Tensor有4种常见的乘法:*, torch.mul, torch.mm, torch.matmul。本文抛砖引玉,简单叙述一下这4种乘法的区别,具体使用还是要参照官方文档。

1、点乘

a与b做*乘法,原则是如果a与b的size不同,则以某种方式将a或b进行复制,使得复制后的a和b的size相同,然后再将a和b做element-wise的乘法。

  • 标量
    Tensor与标量k做*乘法的结果是Tensor的每个元素乘以k(相当于把k复制成与lhs大小相同,元素全为k的Tensor)。
>>> a = torch.ones(3,4)
>>> a
tensor([[1., 1., 1., 1.],
        [1., 1., 1., 1.],
        [1., 1., 1., 1.]])
>>> a * 2
tensor([[2., 2., 2., 2.],
        [2., 2., 2., 2.],
        [2., 2., 2., 2.]])
  • 一维向量
    Tensor与行向量做*乘法的结果是每列乘以行向量对应列的值(相当于把行向量的行复制,成为与lhs维度相同的Tensor). 注意此时要求Tensor的列数与行向量的列数相等。
>>> 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,4])
>>> b
tensor([1., 2., 3., 4.])
>>> a * b
tensor([[1., 2., 3., 4.],
        [1., 2., 3., 4.],
        [1., 2., 3., 4.]])
  • 矩阵
    如果两个二维矩阵A与B做点积A * B,则要求A与B的维度完全相同,即A的行数=B的行数,A的列数=B的列数。
>>> a = torch.tensor([[1, 2], [2, 3]])
>>> a * a
tensor([[1, 4],
        [4, 9]])
  • broadcast
    点积是broadcast的。broadcast是torch的一个概念,简单理解就是在一定的规则下允许高维Tensor和低维Tensor之间的运算。
    这里举一个点积broadcast的例子。在例子中,a是二维Tensor,b是三维Tensor,但是a的维度与b的后两位相同,那么a和b仍然可以做点积,点积结果是一个和b维度一样的三维Tensor,运算规则是:若c = a * b, 则c[i,,] = a * b[i, *, *],即沿着b的第0维做二维Tensor点积,或者可以理解为运算前将a沿着b的第0维也进行了expand操作,即a = a.expand(b.size()); a * b。
>>> a = torch.tensor([[1, 2], [2, 3]])
>>> b = torch.tensor([[[1,2],[2,3]],[[-1,-2],[-2,-3]]])
>>> a * b
tensor([[[ 1,  4],
         [ 4,  9]],

        [[-1, -4],
         [-4, -9]]])
>>> b * a
tensor([[[ 1,  4],
         [ 4,  9]],

        [[-1, -4],
         [-4, -9]]])
2、torch.mul

官方文档关于torch.mul的介绍. 用法与*乘法相同,也是element-wise的乘法,也是支持broadcast的。

3、torch.mm

官方文档关于torch.mm的介绍…数学里的矩阵乘法,要求两个Tensor的维度满足矩阵乘法的要求

例子:

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

官方文档关于torch.matmul的介绍,torch.mm的broadcast版本。

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

        [[4., 4.],
         [4., 4.],
         [4., 4.]],

        [[4., 4.],
         [4., 4.],
         [4., 4.]],

        [[4., 4.],
         [4., 4.],
         [4., 4.]],

        [[4., 4.],
         [4., 4.],
         [4., 4.]]])