Numpy trick：矩阵运算优化

1. 问题场景

按pair-wise计算两组L维向量的平方差距离：

输入矩阵维度为 a：M×L，b：N×L
输出矩阵维度为 c：M×N，其中entry(i,j)为第i个L维a向量和第j个L维b向量间的平方差距离

M, N, L = 10, 20, 50 
a = np.ones((M, L), dtype=np.float32)
b = np.ones((N, L), dtype=np.float32)

2. 解决方案

有两种方案来解决，一种是暴力求解，另一种巧妙利用结构来解决（本质是平方差公式(a-b)^2 = a^2 + b^2 - 2ab）

2.1 方法一

res1 = np.zeros((a.shape[0], b.shape[0]), dtype=np.float32)
for i in range(a.shape[0]):
    for j in range(b.shape[0]):
        for k in range(a.shape[1]):
            res1[i][j] += pow(a[i][k]-b[j][k], 2)

2.2 方法二

这里注意，res2的第一项维度为（M，N），第二项维度为（M，1），第三项维度为（1，N），这三个不同维度的矩阵相加时，会自动进行维度填补，将（M，1）维矩阵按列复制、将（1，N）维矩阵按行复制为（M，N）维。

a2,b2 = np.square(a).sum(axis=1), np.square(b).sum(axis=1)
res2 = -2. * np.dot(a, b.T) + a2[:, None] + b2[None, :]

3. 性能对比

大维度下，运算性能产生巨大的差距！
在这里插入图片描述

4. 完整源码

import numpy as np
from time import time
# 初始化
M, N, L = 200, 300, 500
a = np.ones((M, L), dtype=np.float32)
b = np.ones((N, L), dtype=np.float32)
# 方法一
time1 = time()
res1 = np.zeros((a.shape[0], b.shape[0]), dtype=np.float32)
for i in range(a.shape[0]):
    for j in range(b.shape[0]):
        for k in range(a.shape[1]):
            res1[i][j] += pow(a[i][k]-b[j][k], 2)
time2 = time()
# 方法二
a2,b2 = np.square(a).sum(axis=1), np.square(b).sum(axis=1)
res2 = -2. * np.dot(a, b.T) + a2[:, None] + b2[None, :]
time3 = time()
# 输出
print("维度信息：M="+str(M)+"，N="+str(N)+"，L="+str(L))
print("方法一用时（秒）："+str(time2-time1))
print("方法二用时（秒）："+str(time3-time2))

5. 其他尝试

5.1 更改原矩阵：np.ones -> np.random.rand

在这里插入图片描述

5.2 改变两种方法的执行顺序

M, N, L = 200, 300, 500
# a = np.ones((M,L),dtype=np.float32)
# b = np.ones((N,L),dtype=np.float32)
a = np.random.rand(M, L)
b = np.random.rand(N, L)

time1 = time()
a2,b2 = np.square(a).sum(axis=1), np.square(b).sum(axis=1)
res2 = -2. * np.dot(a, b.T) + a2[:, None] + b2[None, :]
time2 = time()

res1 = np.zeros((a.shape[0], b.shape[0]), dtype=np.float32)
for i in range(a.shape[0]):
    for j in range(b.shape[0]):
        for k in range(a.shape[1]):
            res1[i][j] += pow(a[i][k]-b[j][k], 2)
time3 = time()

print("维度信息：M="+str(M)+"，N="+str(N)+"，L="+str(L))
print("方法二用时（秒）："+str(time2-time1))
print("方法一用时（秒）："+str(time3-time2))