楔子
在numpy中有一個tensordot方法，尤其在做機器學習的時候會很有用。估計有人看到這個名字，會想到tensorflow，沒錯tensorflow里面也有tensordot這個函數。這個函數它的作用就是，可以讓兩個不同維度的數組進行相乘。我們來舉個例子：
import numpy as np
a = np.random.randint(0, 9, (3, 4))
b = np.random.randint(0, 9, (4, 5))
try:
print(a * b)
except Exception as e:
print(e) # operands could not be broadcast together with shapes (3,4) (4,5)
# 很明顯，a和b兩個數組的維度不一樣，沒辦法相乘
# 但是
print(np.tensordot(a, b, 1))
"""
[[32 32 28 28 52]
[10 25 40 38 78]
[56 7 28 0 42]]
"""
# 我們看到使用tensordot是可以的
下面我們來看看這個函數的用法
函數原型
@array_function_dispatch(_tensordot_dispatcher)
def tensordot(a, b, axes=2):
我們看到這個函數接收三個參數，前兩個就是numpy中數組，最后一個參數則是用於指定收縮的軸。它可以接收一個整型、列表、列表里面嵌套列表，具體代表什么含義我們下面舉例說明。
理解axes
axes為整型
如果axes接收的是一個整型：m，那么表示指定數組a的后n個軸和數組b的前n個軸分別進行內積，就是對應位置元素相乘、再整體求和。
import numpy as np
a = np.arange(60).reshape((3, 4, 5))
b = np.arange(160).reshape((4, 5, 8))
# 顯然這兩個數組不能直接相乘，但是a和后兩個軸和b的前兩個軸是可以直接相乘的
# 因為它們都是(4, 5), 最后結果的shape為(3, 8)
print(np.tensordot(a, b, 2).shape) # (3, 8)
而且這個axes默認為2，所以它一般都是針對三維或者三維以上的數組
但是為了具體理解，后面我們會使用一維、二維數據具體舉例說明。現在先看axes取不同的值，會得到什么結果，先理解一下axes的含義。
import numpy as np
a = np.arange(60).reshape((3, 4, 5))
b = np.arange(160).reshape((4, 5, 8))
try:
print(np.tensordot(a, b, 1).shape)
except Exception as e:
print(e) # shape-mismatch for sum
# 結果報錯了，很好理解，就是形狀不匹配嘛
# axes指定為1，表示a的后一個軸和b的前一個軸進行內積
# 但是一個是5一個是4，元素無法一一對應，所以報錯，提示shape-mismatch，形狀不匹配
# 這里我們把數組b的shape改一下，這樣a的后一個軸和b的前一個軸就匹配了，都是5
a = np.arange(60).reshape((3, 4, 5))
b = np.arange(160).reshape((5, 4, 8))
print(np.tensordot(a, b, 1).shape) # (3, 4, 4, 8)
"""
這樣就能夠運算了，我們說指定收縮的軸，進行內積運算得到的是一個值
所以這里的(3, 4, 5)和(5, 4, 8)變成了(3, 4, 4, 8)
而上一個例子是(3, 4, 5)和(4, 5, 8)，然后axes=2
因為a的后兩個軸和b的前兩個軸進行內積變成了一個具體的值，所以最終的維度就是(3, 8)
"""
如果axes為0的話，會有什么結果
import numpy as np
a = np.arange(60).reshape((3, 4, 5))
b = np.arange(160).reshape((4, 5, 8))
print(np.tensordot(a, b, 0).shape) # (3, 4, 5, 4, 5, 8)
print(np.tensordot(b, a, 0).shape) # (4, 5, 8, 3, 4, 5)
"""
np.tensordot(a, b, 0)等價於將a中的每一個元素都和b相乘
然后再將原來a中的對應元素替換掉
"""
上面的操作也可以使用愛因斯坦求和來實現
axes=0
import numpy as np
a = np.arange(60).reshape((3, 4, 5))
b = np.arange(160).reshape((4, 5, 8))
c1 = np.tensordot(a, b, 0)
c2 = np.einsum("ijk,xyz->ijkxyz", a, b)
print(c1.shape, c2.shape) # (3, 4, 5, 4, 5, 8) (3, 4, 5, 4, 5, 8)
print(np.all(c1 == c2)) # True
"""
生成的c1和c2是一樣的
"""
c3 = np.tensordot(b, a, 0)
c4 = np.einsum("ijk,xyz->xyzijk", a, b)
print(c3.shape, c4.shape) # (4, 5, 8, 3, 4, 5) (4, 5, 8, 3, 4, 5)
print(np.all(c3 == c4)) # True
"""
生成的c3和c4是一樣的
"""
那么它們的效率之間孰優孰劣呢？我們在jupyter上測試一下
>>> %timeit c1 = np.tensordot(a, b, 0)
50.5 µs ± 206 ns per loop
>>> %timeit c2 = np.einsum("ijk,xyz->ijkxyz", a, b)
7.29 µs ± 242 ns per loop
可以看到愛因斯坦求和快了不少
axes=1
import numpy as np
a = np.arange(60).reshape((3, 4, 5))
b = np.arange(160).reshape((5, 4, 8))
c1 = np.tensordot(a, b, 1)
c2 = np.einsum("ijk,kyz->ijyz", a, b)
print(c1.shape, c2.shape) # (3, 4, 4, 8) (3, 4, 4, 8)
print(np.all(c1 == c2)) # True
axes=2
import numpy as np
a = np.arange(60).reshape((3, 4, 5))
b = np.arange(160).reshape((4, 5, 8))
c1 = np.tensordot(a, b, 2)
c2 = np.einsum("ijk,jkz->iz", a, b)
print(c1.shape, c2.shape) # (3, 8) (3, 8)
print(np.all(c1 == c2)) # True
axes為列表
如果axes接收的是一個列表：[m, n]，那么表示讓a的第m+1個(索引為m)軸和b的第n+1(索引為n)個軸進行內積。使用列表的方法最大的好處就是，可以指定任意位置的軸。
import numpy as np
a = np.arange(60).reshape((3, 4, 5))
b = np.arange(160).reshape((4, 5, 8))
# 我們看到a的第二個維度(或者說軸)和b的第一個維度都是4，所以它們是可以進行內積的
c1 = np.tensordot(a, b, [1, 0])
# 由於內積的結果是一個標量，所以(3, 4, 5)和(4, 5, 8)在tensordot之后的shape是(3, 5, 5, 8)
# 相當於把各自的4給扔掉了(因為變成了標量)，然后組合在一起
print(c1.shape) # (3, 5, 5, 8)
# 同理a的最后一個維度和b的第二個維度也是可以內積的
# 最后一個維度也可以使用-1，等於按照列表的索引來取對應的維度
c2 = np.tensordot(a, b, [-1, 1])
print(c2.shape) # (3, 4, 4, 8)
上面的操作也可以使用愛因斯坦求和來實現
import numpy as np
a = np.arange(60).reshape((3, 4, 5))
b = np.arange(160).reshape((4, 5, 8))
c1 = np.tensordot(a, b, [1, 0])
c2 = np.einsum("ijk,jyz->ikyz", a, b)
print(c1.shape, c2.shape) # (3, 5, 5, 8) (3, 5, 5, 8)
print(np.all(c1 == c2)) # True
c3 = np.tensordot(a, b, [-1, 1])
c4 = np.einsum("ijk,akz->ijaz", a, b)
print(c3.shape, c4.shape) # (3, 4, 4, 8) (3, 4, 4, 8)
print(np.all(c3 == c4)) # True
axes為列表嵌套列表
如果axes接收的是一個嵌套列表的列表：[[m], [n]]，等於說可以選多個軸
import numpy as np
a = np.arange(60).reshape((3, 4, 5))
b = np.arange(160).reshape((4, 5, 8))
# 我們想讓a的后兩個軸和b的前兩個軸內積
c1 = np.tensordot(a, b, axes=2)
c2 = np.tensordot(a, b, [[1,2], [0,1]])
print(c1.shape, c2.shape) # (3, 8) (3, 8)
print(np.all(c1 == c2)) # True
但是使用列表進行篩選還有一個好處，就是可以忽略順序
import numpy as np
a = np.arange(60).reshape((4, 3, 5))
b = np.arange(160).reshape((4, 5, 8))
# 這個時候就無法給axes傳遞整型了
c3 = np.tensordot(a, b, [[0, 2], [0, 1]])
print(c3.shape) # (3, 8)
此外，使用列表篩選還有一個強大的功能，就是可以倒着取值
import numpy as np
a = np.arange(60).reshape((4, 5, 3))
b = np.arange(160).reshape((5, 4, 8))
# 這個時候我們選擇前兩個軸，但是一個是(4, 5)一個是(5, 4)，所以無法相乘
# 因此在選擇的時候需要倒着篩選:
# [[0, 1], [1, 0]]-> (4, 5)和(4, 5) 或者 [[1, 0], [0, 1]] -> (5, 4)和(5, 4)
c3 = np.tensordot(a, b, [[0, 1], [1, 0]])
print(c3.shape) # (3, 8)
最后同樣看看如何愛因斯坦求和來實現
import numpy as np
a = np.arange(60).reshape((4, 5, 3))
b = np.arange(160).reshape((4, 5, 8))
c1 = np.tensordot(a, b, [[0, 1], [0, 1]])
c2 = np.einsum("ijk,ijz->kz", a, b)
print(c1.shape, c2.shape) # (3, 8) (3, 8)
print(np.all(c1 == c2)) # True
a = np.arange(60).reshape((4, 5, 3))
b = np.arange(160).reshape((5, 4, 8))
c1 = np.tensordot(a, b, [[0, 1], [1, 0]])
c2 = np.einsum("ijk,jiz->kz", a, b)
print(c1.shape, c2.shape) # (3, 8) (3, 8)
print(np.all(c1 == c2)) # True
a = np.arange(60).reshape((4, 3, 5))
b = np.arange(160).reshape((5, 4, 8))
c1 = np.tensordot(a, b, [[0, 2], [1, 0]])
c2 = np.einsum("ijk,kiz->jz", a, b)
print(c1.shape, c2.shape) # (3, 8) (3, 8)
print(np.all(c1 == c2)) # True
以兩個一維數組為例
我們來通過打印具體的數組來看一下tensordot
import numpy as np
a = np.array([1, 2, 3])
b = np.array([2, 3, 4])
print(np.tensordot(a, b, axes=0))
"""
[[ 2 3 4]
[ 4 6 8]
[ 6 9 12]]
"""
print(np.einsum("i,j->ij", a, b))
"""
[[ 2 3 4]
[ 4 6 8]
[ 6 9 12]]
"""
# 我們axes=0，等於是a的每一個元素和相乘，然后再把原來a對應的元素替換掉
# 所以是a中的1 2 3分別和b相乘，得到[2 3 4] [4 6 8] [6 9 12]、再替換掉1 2 3
# 所以結果是[[2 3 4] [4 6 8] [6 9 12]]
如果axes=1呢？
import numpy as np
a = np.array([1, 2, 3])
b = np.array([2, 3, 4])
print(np.tensordot(a, b, axes=1)) # 20
"""
選取a的前一個軸和b的后一個軸進行內積
而a和b只有一個軸，所以結果是一個標量
"""
print(np.einsum("i,i->", a, b)) # 20
如果axes=2呢？首先我們說axes等於一個整型，表示選取a的后n個軸，b的前n個軸，而一維數組它們只有一個軸
import numpy as np
a = np.array([1, 2, 3])
b = np.array([2, 3, 4])
try:
print(np.tensordot(a, b, axes=2)) # 20
except Exception as e:
print(e) # tuple index out of range
顯然索引越界了。
以一個一維數組和一個二維數組為例
我們通過一維數組和二維數組進行tensordot來感受一下
axes=0
import numpy as np
a = np.array([1, 2, 3])
b = np.array([[2, 3, 4]])
print(np.tensordot(a, b, 0))
"""
[[[ 2 3 4]]
[[ 4 6 8]]
[[ 6 9 12]]]
"""
print(np.einsum("i,jk->ijk", a, b))
"""
[[[ 2 3 4]]
[[ 4 6 8]]
[[ 6 9 12]]]
"""
# 很好理解，就是1 2 3分別和[[2, 3, 4]]相乘再替換掉 1 2 3
print(np.tensordot(a, b, 0).shape) # (3, 1, 3)
##########################
print(np.tensordot(b, a, 0))
"""
[[[ 2 4 6]
[ 3 6 9]
[ 4 8 12]]]
"""
print(np.einsum("i,jk->jki", a, b))
"""
[[[ 2 4 6]
[ 3 6 9]
[ 4 8 12]]]
"""
# 很好理解，就是2 3 4分別和[1 2 3]相乘再替換掉 2 3 4
print(np.tensordot(b, a, 0).shape) # (1, 3, 3)
axes=1的話呢？
import numpy as np
a = np.array([1, 2, 3])
b = np.array([[2, 3, 4], [4, 5, 6]])
try:
print(np.tensordot(a, b, 1))
except Exception as e:
print(e) # shape-mismatch for sum
# 我們注意到報錯了，因為axes=1，表示取a的后一個軸和b的前1個軸
# a的shape是(3, 0)，所以它的后一個軸和前一個軸對應的數組長度都是3
# 但是b的前一個軸對應的數組長度是2，不匹配所以報錯
print(np.tensordot(b, a, 1)) # [20 32]
# 我們看到這個是可以的，因為這表示b的后一個軸，數組長度為3，是匹配的
# 讓后一個軸的[2 3 4]、[4 5 6]分別和[1 2 3]進行內積，最終得到兩個標量
try:
print(np.einsum("i,ij->ij", a, b))
except Exception as e:
print(e)
# operands could not be broadcast together with remapped shapes [original->remapped]: (3,)->(3,newaxis) (2,3)->(2,3)
# 同樣對於愛因斯坦求和也是無法這么做的,我們需要換個順序
print(np.einsum("i,ji->j", a, b)) # [20 32]
# 或者
print(np.einsum("j,ij->i", a, b)) # [20 32]
axes=2的話呢？
import numpy as np
a = np.array([1, 2, 3])
b = np.array([[2, 3, 4], [4, 5, 6]])
try:
print(np.tensordot(a, b, 2))
except Exception as e:
print(e) # tuple index out of range
# 我們注意到報錯了，因為axes=2，表示取a的后兩個軸和b的前兩個軸
# 而a總共才1個軸，所以報錯了
try:
print(np.tensordot(b, a, 2))
except Exception as e:
print(e) # shape-mismatch for sum
# 我們看到雖然也報錯了，但是不是報索引越界。
# 因為上面表示取a的前兩個軸，雖然a只有一個，但是此時不會索引越界，只是就取一個。如果是取后兩個就會越界了
# 此時b是(2, 3)，而a是(3,) 不匹配，可能有人覺得會發生廣播，但在這里不會
以兩個二維數組為例
我們再通過兩個二維數組進行tensordot來感受一下
axes=0
import numpy as np
a = np.array([[1, 2, 3]])
b = np.array([[2, 3, 4], [4, 5, 6]])
# a_shape: (1, 3) b_shape(3, 3)
print(np.tensordot(a, b, 0))
"""
[[[[ 2 3 4]
[ 4 5 6]]
[[ 4 6 8]
[ 8 10 12]]
[[ 6 9 12]
[12 15 18]]]]
"""
print(np.einsum("ij,xy->ijxy", a, b))
"""
[[[[ 2 3 4]
[ 4 5 6]]
[[ 4 6 8]
[ 8 10 12]]
[[ 6 9 12]
[12 15 18]]]]
"""
print(np.tensordot(a, b, 0).shape) # (1, 3, 2, 3)
#############
print(np.tensordot(b, a, 0))
"""
[[[[ 2 4 6]]
[[ 3 6 9]]
[[ 4 8 12]]]
[[[ 4 8 12]]
[[ 5 10 15]]
[[ 6 12 18]]]]
"""
print(np.einsum("ij,xy->xyij", a, b))
"""
[[[[ 2 4 6]]
[[ 3 6 9]]
[[ 4 8 12]]]
[[[ 4 8 12]]
[[ 5 10 15]]
[[ 6 12 18]]]]
"""
print(np.tensordot(b, a, 0).shape) # (2, 3, 1, 3)
axes=1
import numpy as np
a = np.array([[1, 2], [3, 4]])
b = np.array([[2, 3, 4], [4, 5, 6]])
# a_shape: (2, 2) b_shape(2, 3)
print(np.tensordot(a, b, 1))
"""
[[10 13 16]
[22 29 36]]
"""
print(np.einsum("ij,jk->ik", a, b))
"""
[[10 13 16]
[22 29 36]]
"""
# 仔細的你肯定發現了，此時就相當於矩陣的點乘
print(a @ b)
"""
[[10 13 16]
[22 29 36]]
"""
axes=2
import numpy as np
a = np.array([[1, 2], [3, 4]])
b = np.array([[2, 3, 4], [4, 5, 6]])
# a_shape: (2, 2) b_shape(2, 3)
# 取后兩個軸顯然不行，因為(2, 2)和(2, 3)不匹配
try:
print(np.tensordot(a, b, 2))
except Exception as e:
print(e) # shape-mismatch for sum
a = np.array([[1, 2, 3], [2, 2, 2]])
b = np.array([[2, 3, 4], [4, 5, 6]])
print(np.tensordot(a, b, 2)) # 50
print(np.einsum("ij,ij->", a, b)) # 50
最后看即個愛因斯坦求和的例子，感受它和主角tensordot的區別，當然如果不熟悉的愛因斯坦求和的話可以不用看
import numpy as np
a = np.random.randint(1, 9, (5, 3, 2, 3))
b = np.random.randint(1, 9, (3, 3, 2))
c1 = a @ b # 多維數組，默認是對最后兩位進行點乘
c2 = np.einsum("ijkm,jmn->ijkn", a, b)
print(np.all(c1 == c2)) # True
print(c2.shape) # (5, 3, 2, 2)
print(np.einsum("...km,...mn->...kn", a, b).shape) # (5, 3, 2, 2)
# 但如果是
c3 = np.einsum("ijkm,amn->ijkn", a, b)
print(c3.shape) # (5, 3, 2, 2)
# 由於符號不一樣，所以即使shape一致，但是兩個數組不一樣
print(np.all(c3 == c1)) # False
a = np.random.randint(1, 9, (5, 3, 3, 2))
b = np.random.randint(1, 9, (1, 3, 2))
print(np.einsum("ijmk,jmn->ijkn", a, b).shape) # (5, 3, 2, 2)
print(np.einsum("ijkm,jnm->ijkn", a, b).shape) # (5, 3, 3, 3)