许多数值计算库都具有高效的矢量化操作实现。矩阵乘法、查找点积等操作非常高效。这些操作的实现是为了利用 CPU 中的多个核心,并将计算卸载到 GPU(如果可用)。通常,矩阵和向量的操作由 BLAS(基本线性代数子程序)提供。一些示例是 Intel MKL、OpenBLAS、cuBLAS 等。
首先我们创建两个矩阵,以便我们可以用它来检查我们的实现是否正确。
import tensorflow as tfimport numpy as np tf.__version__# 2.0.0 a = np.random.normal(size=(200, 784)).astype('float32') b = np.random.normal(size=(784, 10)).astype('float32') expected = np.matmul(a, b)
不使用外部库实现功能,
defpy_matmul1(a,b): ra, ca = a.shape rb, cb = b.shapeassert ca == rb,f"{ca} != {rb}" output = np.zeros(shape=(ra, cb))for i inrange(ra):for j inrange(cb):for k inrange(rb): output[i, j]+= a[i, k]* b[k, j]return output %time result =py_matmul1(a, b)assert result.shape == expected.shapeassert np.allclose(result, expected, rtol=1e-02), (result, expected)
执行时,在我的计算机上需要 1.38 秒。它相当慢,可以大大改进。如果您注意到,最内层循环基本上是在计算两个向量的点积。在这种情况下,两个向量分别是 a 和 b 的第 i 行和第 j 列。所以让我们用点积实现删除最内层循环。
矢量优化一:
defpy_matmul2(a,b): ra, ca = a.shape rb, cb = b.shapeassert ca == rb,f"{ca} != {rb}" output = np.zeros(shape=(ra, cb))for i inrange(ra):for j inrange(cb):# we replaced the loop with dot product output[i, j]= np.dot(a[i], b[:,j])return output %time result =py_matmul2(a, b)assert result.shape == expected.shapeassert np.allclose(result, expected, rtol=1e-02), (result, expected)
此实现仅需 6 毫秒。与原始实现相比,这是一个巨大的改进。由于内部循环本质上是在计算点积,我们将其替换为 np.dot 函数,并传递矩阵 a 中的第 i 行和矩阵 b 中的第 j 列。
矢量优化二:
现在让我们删除迭代矩阵 b 的列的 for 循环。
defpy_matmul3(a,b): ra, ca = a.shape rb, cb = b.shapeassert ca == rb,f"{ca} != {rb}" output = np.zeros(shape=(ra, cb))for i inrange(ra): output[i]= np.dot(a[i], b)return output %time result =py_matmul3(a, b)assert result.shape == expected.shapeassert np.allclose(result, expected, rtol=1e-02), (result, expected)
此实现耗时 2.97 毫秒。使用称为广播的技术,我们可以从本质上消除循环,仅使用一行 output[i] = np.dot(a[i], b),我们就可以计算输出矩阵第 i 行的整个值。numpy 所做的是广播向量 a[i],使其与矩阵 b 的形状相匹配。然后它计算每对向量的点积。广播规则在 numpy、tensorflow、pytorch 等主要库中几乎相同。
库优化三:
现在让我们使用 numpy 的内置 matmul 函数。
defpy_matmul4(a,b): ra, ca = a.shape rb, cb = b.shapeassert ca == rb,f"{ca} != {rb}"return np.matmul(a, b) %time result =py_matmul4(a, b)assert result.shape == expected.shapeassert np.allclose(result, expected, rtol=1e-02), (result, expected)
使用numpy的内置matmul函数,需要999 μ 𝜇 s。这是迄今为止我们实施的最快的。
库优化四:
在张量流中,它也与 numpy 非常相似。我们只需要调用 matmul 函数即可。
defpy_matmul5(a,b): ra, ca = a.shape rb, cb = b.shapeassert ca == rb,f"{ca} != {rb}"return tf.matmul(a, b) tf_a = tf.constant(a) tf_b = tf.constant(b)%time result =py_matmul5(tf_a, tf_b)assert result.shape == expected.shapeassert np.allclose(result, expected, rtol=1e-02), (result, expected)