從手寫三層循環到標準實現，矩陣相乘運行效率提升三萬六千倍之路

前言

矩陣乘法能夠說是最多見的運算之一。

本文介紹不一樣的方式實現的矩陣乘法，並比較它們運行速度的差別。

表示矩陣的方式有不少種，完善的矩陣類應該實現切片取值，得到矩陣形狀等操做，但本文並不打算直接從原生Python實現一個矩陣類，而是直接用 Pytorch中的tensor表示矩陣。

開始: 三層循環

根據矩陣相乘定義，可經過三層循環實現該運算。

def matmul(a, b):
    r1, c1 = a.shape
    r2, c2 = b.shape
    
    assert c1 == r2
    
    rst = torch.zeros(r1, c2)
    
    for i in range(r1):
        for j in range(c2):
            for k in range(c1):
                rst[i][j] += a[i][k] * b[k][j]
    return rst

那麼這個函數的運行效率如何呢？讓咱們嘗試兩個較大的矩陣相乘，測試一下運行時間。

m1 = torch.randn(5, 784)
m2 = torch.randn(784, 10)

%timeit -n 10 matmul(m1, m2)

獲得結果以下：

624 ms ± 3.76 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

即每次矩陣相乘， 須要耗時 600ms 左右，這是一個很是很是慢的速度，慢到兩次矩陣乘法竟然要耗時1秒多，這是不可能被接受的。

相同形狀的張量進行運算

若是兩個張量的形狀相同，則他們的運算爲同一位置的數字進行運算。

a = torch.tensor([1., 2, 3])
b = torch.tensor([4., 5, 6])

a + b  # tensor([5., 7., 9.])
a * b  # tensor([ 4., 10., 18.])

康康以前用三層循環實現的矩陣相乘，發現最裏面一層循環的本質就是兩個一樣大小的張量相乘，再進行求和。
即第一個矩陣中的一行 跟 第二個矩陣中的一列 進行運算，且這行和列中的元素個數相同，則咱們能夠經過一樣形狀的張量運算改寫最內層循環：

def matmul(a, b):
    r1, c1 = a.shape
    r2, c2 = b.shape
    
    assert c1 == r2
    
    rst = torch.zeros(r1, c2)
    
    for i in range(r1):
        for j in range(c2):
            rst[i][j] = (a[i,:] * b[:,j]).sum()  # 改了這裏
    return rst

%timeit -n 10 matmul(m1, m2)

獲得結果以下

1.4 ms ± 92.2 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

624 / 1.4=445，只改寫了一下最內層循環，就使得矩陣乘法快了445倍！

廣播機制

廣播機制使得不一樣形狀的張量間能夠進行運算：

1. 兩個張量擴充成一樣的形狀
2. 再按相同形狀的張量進行運算

# shape: [2, 3]
a = torch.tensor([
    [1, 2, 3],
    [4, 5, 6],
])

# shape: [1]
b = torch.tensor([1])

# shape: [3]
c = torch.tensor([10, 20, 30])

形狀爲 [2, 3] 和 [1] 的兩個張量相加:

a + b

"""輸出:
tensor([[2, 3, 4],
        [5, 6, 7]])
"""

形狀爲 [2, 3] 和 [3] 的兩個張量相加:

b + c

"""輸出:
tensor([[11, 22, 33],
        [14, 25, 36]])
"""

這兩個例子中，維度低的張量都是暗地裏先擴充成了維度高的張量，而後再參與的運算。

那麼如何查看擴充後的張量是啥呢？用 expand_as 函數就能夠查看：

b.expand_as(a)

"""輸出
tensor([[1, 1, 1],
        [1, 1, 1]])
"""

b.expand_as(a)

"""輸出
tensor([[10, 20, 30],
        [10, 20, 30]])
"""

這就一目瞭然了，形狀不一樣的張量能夠經過廣播機制擴充成形狀一致的張量再進行運算。

那麼任意形狀的兩個張量均可以運算嗎？固然不是了，判斷兩個張量是否能運算的規則以下：

先從兩個張量的最後一個維度看起，若是維度的維數相同，或者其中一個維數爲1，則能夠繼續判斷，不然就失敗。
而後看倒數第二個維度，倒數第三個維數，一直到遍歷完某個張量的維數爲止，一直沒有失敗則這兩個張量能夠經過廣播機制進行運算。

那麼這個廣播機制和矩陣乘法有什麼關係呢？答案就是它能夠幫咱們再去掉一層循環。

如今的最內存循環的本質是 一個形狀爲 [c1] 的張量 和 一個形狀爲 [c1, c2] 的張量作運算，最終生成一個形狀爲 [c2] 的張量。

則咱們能夠把矩陣運算改寫爲:

def matmul(a, b):
    r1, c1 = a.shape
    r2, c2 = b.shape
    
    assert c1 == r2
    
    rst = torch.zeros(r1, c2)
    
    for i in range(r1):
        rst[i] = (a[i, :].unsqueeze(-1) * b).sum(0)
    return rst

%timeit -n 10 matmul(m1, m2)

"""輸出
249 µs ± 66.4 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
"""

如今已經把每次矩陣運算的時間壓縮到了 249 µs！！！，比最開始的 624ms 快了 2500倍！

對於 unsqueeze 操做不太熟悉的小夥伴請看個人另外一篇文檔: Pytorch 中張量的理解

可是還沒結束。。。由於兩個矩陣的相乘，就是 [r1, c1] 和 [c1, c2] 兩個張量的運算，咱們能夠直接把它用廣播機制一次到位的算出結果，連惟一的那層循環也能夠省去：

def matmul(a, b):
    r1, c1 = a.shape
    r2, c2 = b.shape
    
    assert c1 == r2
    
    return (a.unsqueeze(-1) * b.unsqueeze(0)).sum(1)

%timeit -n 10 matmul(m1, m2)

"""輸出:
169 µs ± 41.6 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
"""

這個 169µs 已是最開始矩陣相乘版本的 3700 倍了。。( Ĭ ^ Ĭ )淚目，果真知識是第一輩子產力。

愛因斯坦求和

接下來就是 pytorch 自帶的矩陣運算工具了，其中一個是愛因斯坦求和，貌似知道這個的同窗很少。。
簡單來講，它能讓咱們幾乎不編寫代碼就能進行矩陣運算，只須要肯定輸入和輸出矩陣的形狀便可：

def matmul(a, b):
    return torch.einsum("ik,kj->ij", a, b)

%timeit -n 10 matmul(a, b)

"""輸出
74 µs ± 25.6 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
"""

74µs 這個速度已是原始版本的 8000 多倍了。。。可是對於工業級別的要求彷佛仍然不夠快~

pytorch 的矩陣相乘標準實現

最後祭出 pytorch 的矩陣相乘官方版本：

def matmul(a, b):
    return a @ b

%timeit -n 10 matmul(m1, m2)

"""輸出
17.1 µs ± 28.5 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
"""

17.1 µs 是原始三層循環版本的 36000 倍，官方實現就是這麼簡單枯燥，樸實無華~