【tensorflow】多维张量做tf.matmul

首发地址：https://zhuanlan.zhihu.com/p/138731311
线性代数都学过二维矩阵的乘法，而tf.matmul还可以处理多维矩阵，比如

import tensorflow as tf
import numpy as np
a = tf.random.uniform([2, 1, 2, 3])
b = tf.random.uniform([1, 3, 3, 2])
c = tf.matmul(a, b)

c是什么呢？

先给出结论：不管几维矩阵都是先做最后两维的矩阵的乘法，再在不同维度重复多次。

多维的 tf.matmul(a, b) 的维度有如下两个要求：

1、a的axis=-1的值（_只可意会）和b的axis=-2的值需要相等。比如上述例子中[3, 2, 3]最后的3，和[3, 3, 2]的第二个3。

2、a和b的各维度的值（除了axis=-1和-2的值），在任意维度上，都需要“相等”或“有一个是1”。

比如，[3, 2, 3]维度的张量与[3, 3, 2]维度的张量做tf.matmul的例子：

In [84]: import tensorflow as tf
    ...: import numpy as np
    ...: a = tf.random.uniform([3, 2, 3])
    ...: b = tf.random.uniform([3, 3, 2])
    ...: c = tf.matmul(a, b)
    ...: c.shape
    ...:
    ...:

Out[84]: TensorShape([3, 2, 2])

In [87]: tf.matmul(a[0],b[0])
Out[87]:
<tf.Tensor: id=374, shape=(2, 2), dtype=float32, numpy=
array([[1.4506222 , 1.323427  ],
       [0.28268352, 0.2917934 ]], dtype=float32)>

In [88]: tf.matmul(a[1],b[1])
Out[88]:
<tf.Tensor: id=383, shape=(2, 2), dtype=float32, numpy=
array([[1.0278544 , 0.4219831 ],
       [0.865297  , 0.87740964]], dtype=float32)>

In [89]: c
Out[89]:
<tf.Tensor: id=365, shape=(3, 2, 2), dtype=float32, numpy=
array([[[1.4506222 , 1.323427  ],
        [0.28268352, 0.2917934 ]],

       [[1.0278544 , 0.4219831 ],
        [0.865297  , 0.8774096 ]],

       [[0.5752927 , 0.13066964],
        [0.5343988 , 0.2741483 ]]], dtype=float32)>

可以看到，[3, 2, 3]维度的张量与[3, 3, 2]维度的张量做tf.matmul，可以理解成：

第一步，先在axis=1和2的维度上做[2, 3]维度的张量与[3, 2]维度的张量之间的二维张量的矩阵乘法，得到[2, 2]维度的结果；

第二部，然后在axis=0的维度上，分别选a的第i个和选b的第i个做上述的第一步，最终得到[3, 2，2]维度的输出。

如果，a和b的axis=0维度对不上，会bug：

In [95]: import tensorflow as tf
    ...: import numpy as np
    ...: a = tf.random.uniform([2, 2, 3])
    ...: b = tf.random.uniform([3, 3, 2])
    ...: c = tf.matmul(a, b)
    ...: c.shape
    ...:
    ...:
---------------------------------------------------------------------------
InvalidArgumentError                      Traceback (most recent call last)
<ipython-input-95-462c4976a35a> in <module>
      3 a = tf.random.uniform([2, 2, 3])
      4 b = tf.random.uniform([3, 3, 2])
----> 5 c = tf.matmul(a, b)
      6 c.shape
      7

D:\S\Anaconda3_v3\lib\site-packages\tensorflow_core\python\util\dispatch.py in wrapper(*args, **kwargs)
    178     """Call target, and fall back on dispatchers if there is a TypeError."""
    179     try:
--> 180       return target(*args, **kwargs)
    181     except (TypeError, ValueError):
    182       # Note: convert_to_eager_tensor currently raises a ValueError, not a

D:\S\Anaconda3_v3\lib\site-packages\tensorflow_core\python\ops\math_ops.py in matmul(a, b, transpose_a, transpose_b, adjoint_a, adjoint_b, a_is_sparse, b_is_sparse, name)
   2725         b = conj(b)
   2726         adjoint_b = True
-> 2727       return batch_mat_mul_fn(a, b, adj_x=adjoint_a, adj_y=adjoint_b, name=name)
   2728
   2729     # Neither matmul nor sparse_matmul support adjoint, so we conjugate

D:\S\Anaconda3_v3\lib\site-packages\tensorflow_core\python\ops\gen_math_ops.py in batch_mat_mul_v2(x, y, adj_x, adj_y, name)
   1700       else:
   1701         message = e.message
-> 1702       _six.raise_from(_core._status_to_exception(e.code, message), None)
   1703   # Add nodes to the TensorFlow graph.
   1704   if adj_x is None:

D:\S\Anaconda3_v3\lib\site-packages\six.py in raise_from(value, from_value)

InvalidArgumentError: In[0] and In[1] must have compatible batch dimensions: [2,2,3] vs. [3,3,2] [Op:BatchMatMulV2] name: MatMul/       

但是当a和b中axis=0的值有一个是1，不会bug：

In [90]: import tensorflow as tf
    ...: import numpy as np
    ...: a = tf.random.uniform([1, 2, 3])
    ...: b = tf.random.uniform([3, 3, 2])
    ...: c = tf.matmul(a, b)
    ...: c.shape
    ...:
    ...:
Out[90]: TensorShape([3, 2, 2])

In [91]: c
Out[91]:
<tf.Tensor: id=398, shape=(3, 2, 2), dtype=float32, numpy=
array([[[0.59542704, 0.60751694],
        [0.19115494, 0.36344892]],

       [[1.0542538 , 0.75257593],
        [0.26940605, 0.24408351]],

       [[1.1716111 , 0.4058628 ],
        [0.09086016, 0.28043625]]], dtype=float32)>

In [92]: tf.matmul(a[0],b[0])
Out[92]:
<tf.Tensor: id=407, shape=(2, 2), dtype=float32, numpy=
array([[0.59542704, 0.60751694],
       [0.19115494, 0.36344892]], dtype=float32)>

In [93]: tf.matmul(a[0],b[1])
Out[93]:
<tf.Tensor: id=416, shape=(2, 2), dtype=float32, numpy=
array([[1.0542538 , 0.7525759 ],
       [0.26940605, 0.2440835 ]], dtype=float32)>

In [94]: tf.matmul(a[0],b[2])
Out[94]:
<tf.Tensor: id=425, shape=(2, 2), dtype=float32, numpy=
array([[1.1716112 , 0.4058628 ],
       [0.09086016, 0.28043625]], dtype=float32)>

依然遵循上述的先最后两维做乘法，再依次组成结果，只是由于a的axis=0的值为1，所以是b在axis=0的所有都对应a的axis=0的唯一。（还是看代码和输出结果更清楚。）

所以得到三维上的结论：

先做最后两维的矩阵的乘法，再在不同维度重复多次。

多维的 tf.matmul(a, b) 的维度有如下两个要求：

1、a的axis=2的值（_只可意会）和b的axis=1的值需要相等。

2、a和b的axis=0的值需要“相等”或者“有一个是1”。

再看更高维度，比如四维的情况。

In [96]: import tensorflow as tf
    ...: import numpy as np
    ...: a = tf.random.uniform([2, 1, 2, 3])
    ...: b = tf.random.uniform([2, 3, 3, 2])
    ...: c = tf.matmul(a, b)
    ...: c.shape
    ...:
    ...:
Out[96]: TensorShape([2, 3, 2, 2])

In [97]: c
Out[97]:
<tf.Tensor: id=454, shape=(2, 3, 2, 2), dtype=float32, numpy=
array([[[[1.0685383 , 1.9015994 ],
         [1.1457413 , 1.5246255 ]],

        [[0.953201  , 1.5544493 ],
         [0.7639411 , 1.4360913 ]],

        [[0.67427766, 0.49847895],
         [0.499685  , 0.39281937]]],


       [[[0.42752475, 0.7453967 ],
         [0.3735991 , 0.74812794]],

        [[0.54442215, 0.6510606 ],
         [0.6632798 , 0.38497943]],

        [[0.3459217 , 0.96300673],
         [0.45035997, 0.90772474]]]], dtype=float32)>

In [98]: tf.matmul(a[0],b[0])
Out[98]:
<tf.Tensor: id=463, shape=(3, 2, 2), dtype=float32, numpy=
array([[[1.0685383 , 1.9015994 ],
        [1.1457413 , 1.5246255 ]],

       [[0.953201  , 1.5544493 ],
        [0.7639411 , 1.4360913 ]],

       [[0.67427766, 0.49847895],
        [0.499685  , 0.39281937]]], dtype=float32)>

In [99]: tf.matmul(a[1],b[1])
Out[99]:
<tf.Tensor: id=472, shape=(3, 2, 2), dtype=float32, numpy=
array([[[0.42752475, 0.7453967 ],
        [0.3735991 , 0.74812794]],

       [[0.54442215, 0.6510606 ],
        [0.6632798 , 0.38497943]],

       [[0.3459217 , 0.96300673],
        [0.45035997, 0.90772474]]], dtype=float32)>

和三维时候是一致的，层层都依次做tf.matmul，也都能转化为最后两维的二维矩阵乘法。

同理，axis=0维度位置的值，有一个是1，也行：

In [100]: import tensorflow as tf
     ...: import numpy as np
     ...: a = tf.random.uniform([2, 1, 2, 3])
     ...: b = tf.random.uniform([1, 3, 3, 2])
     ...: c = tf.matmul(a, b)
     ...: c.shape
     ...:
     ...:
Out[100]: TensorShape([2, 3, 2, 2])

不再赘述

最终结论：不管几维矩阵都是先做最后两维的矩阵的乘法，再在不同维度重复多次。

多维的 tf.matmul(a, b) 的维度有如下两个要求：

1、a的axis=-1的值（_只可意会）和b的axis=-2的值需要相等。

2、a和b的各维度的值（除了axis=-1和-2的值），在任意维度上，都需要“相等”或“有一个是1”。

另外给出一些维度数量对不上的例子，供意会：

In [105]: import tensorflow as tf
     ...: import numpy as np
     ...: a = tf.random.uniform([2, 1, 2, 3])
     ...: b = tf.random.uniform([1, 3, 2])
     ...: c = tf.matmul(a, b)
     ...: c.shape
Out[105]: TensorShape([2, 1, 2, 2])

In [106]: import tensorflow as tf
     ...: import numpy as np
     ...: a = tf.random.uniform([2, 1, 2, 3])
     ...: b = tf.random.uniform([7, 3, 2])
     ...: c = tf.matmul(a, b)
     ...: c.shape
Out[106]: TensorShape([2, 7, 2, 2])

In [107]: import tensorflow as tf
     ...: import numpy as np
     ...: a = tf.random.uniform([2, 1, 2, 3])
     ...: b = tf.random.uniform([7, 9, 3, 2])
     ...: c = tf.matmul(a, b)
     ...: c.shape
---------------------------------------------------------------------------
InvalidArgumentError                      Traceback (most recent call last)
<ipython-input-107-ff6e40117cf7> in <module>
      3 a = tf.random.uniform([2, 1, 2, 3])
      4 b = tf.random.uniform([7, 9, 3, 2])
----> 5 c = tf.matmul(a, b)
      6 c.shape

D:\S\Anaconda3_v3\lib\site-packages\tensorflow_core\python\util\dispatch.py in wrapper(*args, **kwargs)
    178     """Call target, and fall back on dispatchers if there is a TypeError."""
    179     try:
--> 180       return target(*args, **kwargs)
    181     except (TypeError, ValueError):
    182       # Note: convert_to_eager_tensor currently raises a ValueError, not a

D:\S\Anaconda3_v3\lib\site-packages\tensorflow_core\python\ops\math_ops.py in matmul(a, b, transpose_a, transpose_b, adjoint_a, adjoint_b, a_is_sparse, b_is_sparse, name)
   2725         b = conj(b)
   2726         adjoint_b = True
-> 2727       return batch_mat_mul_fn(a, b, adj_x=adjoint_a, adj_y=adjoint_b, name=name)
   2728
   2729     # Neither matmul nor sparse_matmul support adjoint, so we conjugate

D:\S\Anaconda3_v3\lib\site-packages\tensorflow_core\python\ops\gen_math_ops.py in batch_mat_mul_v2(x, y, adj_x, adj_y, name)
   1700       else:
   1701         message = e.message
-> 1702       _six.raise_from(_core._status_to_exception(e.code, message), None)
   1703   # Add nodes to the TensorFlow graph.
   1704   if adj_x is None:

D:\S\Anaconda3_v3\lib\site-packages\six.py in raise_from(value, from_value)

InvalidArgumentError: In[0] and In[1] must have compatible batch dimensions: [2,1,2,3] vs. [7,9,3,2] [Op:BatchMatMulV2] name: MatMul/

a和b的维度对不上也可以用，规则是“向右看齐”。

后面讨论多维 tf.matmul(a, b, transpose_b=True) 的情况：

In [111]: import tensorflow as tf
     ...: import numpy as np
     ...: a = tf.random.uniform([2, 1, 2, 3])
     ...: b = tf.random.uniform([2, 1, 2, 3])
     ...: c = tf.matmul(a, b, transpose_b=True)
     ...: c.shape
Out[111]: TensorShape([2, 1, 2, 2])

In [112]: import tensorflow as tf
     ...: import numpy as np
     ...: a = tf.random.uniform([2, 1, 2, 3])
     ...: b = tf.random.uniform([1, 5, 2, 3])
     ...: c = tf.matmul(a, b, transpose_b=True)
     ...: c.shape
Out[112]: TensorShape([2, 5, 2, 2])

transpose只是对最后两维做了转置，用于二维矩阵乘法能对的上。

---

感觉有用请点赞~~~谢谢