为什么用 Numpy 还是慢, 你用对了吗?

2,633 阅读11分钟
原文链接: zhuanlan.zhihu.com

最近在写代码, 编一个 Python 模拟器, 做 simulation, 好不容易用传说中 Python 里速度最快的计算模块 "Numpy" 的写好了, 结果运行起来, 出奇的慢! 因为一次simulation要一个小时, 要不停测试, 所以自己受不了了.. 首先, 我的脑海中的问题, 渐渐浮现出来.

  • 我知道 Pandas 要比 Numpy 慢, 所以我尽量避免用 Pandas. 但是 Numpy (速度怪兽), 为什么还是这么慢?

带有写代码洁癖的我好好给 google 了一番. 第一个出现在我眼前的就是这个文章, Getting the Best Performance out of NumPy. 所以我也将自己从这个文章中学到的诀窍分享给大家, 并补充一些内容.

为什么用 Numpy?

我们都知道, Python 是慢的, 简单来说, 因为 Python 执行你代码的时候会执行很多复杂的 "check" 功能, 比如当你赋值

b=1; a=b/0.5

这个运算看似简单, 但是在计算机内部, b 首先要从一个整数 integer 转换成浮点数 float, 才能进行后面的 `b/0.5`, 因为得到的要是一个小数. 还有很多其他的原因和详细说明 (比如 Python 对内存的调用) 在这里能够找到: Why Python is Slow: Looking Under the Hood

提到 Numpy, 它就是一个 Python 的救星. 能把简单好用的 Python 和高性能的 C 语言合并在一起. 当你调用 Numpy 功能的时候, 他其实调用了很多 C 语言而不是纯 Python. 这就是为什么大家都爱用 Numpy 的原因.

创建 Numpy Array 的结构

其实 Numpy 就是 C 的逻辑, 创建存储容器 "Array" 的时候是寻找内存上的一连串区域来存放, 而 Python 存放的时候则是不连续的区域, 这使得 Python 在索引这个容器里的数据时不是那么有效率. Numpy 只需要再这块固定的连续区域前后走走就能不费吹灰之力拿到数据. 下图是来自 Why Python is Slow: Looking Under the Hood, 他很好的解释了这一切.

在运用 Numpy 的时候, 我们通常不是用一个一维 Array 来存放数据, 而是用二维或者三维的块来存放 (说出了学机器学习的朋友们的心声~).

因为 Numpy 快速的矩阵相乘运算, 能将乘法运算分配到计算机中的多个核, 让运算并行. 这年头, 我们什么都想多线程/多进程 (再次说出了机器学习同学们的心声~). 这也是 Numpy 为什么受人喜欢的一个原因. 这种并行运算大大加速了运算速度.

那么对于这种天天要用到的2D/3D Array, 我们通常都不会想着他是怎么来的. 因为按照我们正常人的想法, 这矩阵就是矩阵, 没什么深度的东西呀. 不过这可不然! 要不然我也不会写这篇分享了. 重点来了, 不管是1D/2D/3D 的 Array, 从根本上, 它都是一个 1D array!

这篇 Blog的图显示. 在我们看来的 2D Array, 如果追溯到计算机内存里, 它其实是储存在一个连续空间上的. 而对于这个连续空间, 我们如果创建 Array 的方式不同, 在这个连续空间上的排列顺序也有不同. 这将影响之后所有的事情! 我们后面会用 Python 进行运算时间测试.

在 Numpy 中, 创建 2D Array 的默认方式是 "C-type" 以 row 为主在内存中排列, 而如果是 "Fortran" 的方式创建的, 就是以 column 为主在内存中排列.

col_major = np.zeros((10,10), order='C')    # C-type
row_major = np.zeros((10,10), order='F')    # Fortran

在 axis 上的动作

当你的计算中涉及合并矩阵, 不同形式的矩阵创建方式会给你不同的时间效果. 因为在 Numpy 中的矩阵合并等, 都是发生在一维空间里, ! 不是我们想象的二维空间中!

a = np.zeros((200, 200), order='C')
b = np.zeros((200, 200), order='F')
N = 9999

def f1(a):
    for _ in range(N):
        np.concatenate((a, a), axis=0)

def f2(b):
    for _ in range(N):
        np.concatenate((b, b), axis=0)

t0 = time.time()
f1(a)
t1 = time.time()
f2(b)
t2 = time.time()

print((t1-t0)/N)     # 0.000040
print((t2-t1)/N)     # 0.000070

从上面的那张图, 可以想到, row 为主的存储方式, 如果在 row 的方向上合并矩阵, 将会更快. 因为只要我们将思维放在 1D array 那, 直接再加一个 row 放在1D array 后面就好了, 所以在上面的测试中, f1 速度要更快. 但是在以 column 为主的系统中, 往 1D array 后面加 row 的规则变复杂了, 消耗的时间也变长. 如果以 axis=1 的方式合并, "F" 方式的 f2 将会比 "C" 方式的 f1 更好.

还有一个要提的事情, 为了图方便, 有时候我会直接使用 `np.stack` 来代替 `np.concatenate`, 因为这样可以少写一点代码, 不过使用上面的形式, 通过上面的测试发现是这样. 所以之后为了速度, 我推荐还是尽量使用 `np.concatenate`.

np.vstack((a,a))                # 0.000063
np.concatenate((a,a), axis=0)   # 0.000040

或者有时候在某个 axis 上进行操作, 比如对上面用 "C-type" 创建的 a 矩阵选点:

indices = np.random.randint(0, 100, size=10, dtype=np.int32)
a[indices, :]     # 0.000003
a[:, indices]     # 0.000006

因为 a 是用 row 为主的形式储存, 所以在 row 上面选数据要比在 column 上选快很多! 对于其他的 axis 的操作, 结果也类似. 所以你现在懂了吧, 看自己要在哪个 axis 上动的手脚多, 然后再创建合适于自己的矩阵形式 ("C-type"/"Fortran").

copy慢 view快

在 Numpy 中, 有两个很重要的概念, copy 和 view. copy 顾名思义, 会将数据 copy 出来存放在内存中另一个地方, 而 view 则是不 copy 数据, 直接取源数据的索引部分. 下图来自 Understanding SettingwithCopyWarning in pandas

上面说的是什么意思呢? 我们直接看代码.

a = np.arange(1, 7).reshape((3,2))
a_view = a[:2]
a_copy = a[:2].copy()

a_copy[1,1] = 0
print(a)
"""
[[1 2]
 [3 4]
 [5 6]]
"""

a_view[1,1] = 0
print(a)
"""
[[1 2]
 [3 0]
 [5 6]]
"""

简单说, a_view 的东西全部都是 a 的东西, 动 a_view 的任何地方, a 都会被动到, 因为他们在内存中的位置是一模一样的, 本质上就是自己. 而 a_copy 则是将 a copy 了一份, 然后把 a_copy 放在内存中的另外的地方, 这样改变 a_copy, a 是不会被改变的.

那为什么要提这点呢? 因为 view 不会复制东西, 速度快! 我们来测试一下速度. 下面的例子中 `a*=2` 就是将这个 view 给赋值了, 和 `a[:] *= 2` 一个意思, 从头到尾没有创建新的东西. 而 `b = 2*b` 中, 我们将 b 赋值给另外一个新建的 b.

a = np.zeros((1000, 1000))
b = np.zeros((1000, 1000))
N = 9999

def f1(a):                 
    for _ in range(N):
        a *= 2           # same as a[:] *= 2

def f2(b):
    for _ in range(N):
        b = 2*b

print('%f' % ((t1-t0)/N))     # f1: 0.000837
print('%f' % ((t2-t1)/N))     # f2: 0.001346

对于 view 还有一点要提, 你是不是偶尔有时候要把一个矩阵展平, 用到 `np.flatten()` 或者 `np.ravel()`. 他俩是不同的! ravel 返回的是一个 view (谢谢评论中 @非易 的提醒, 官方说如果用 ravel, 需要 copy 的时候才会被 copy , 我想这个时候可能是把 ravel 里面 order 转换的时候, 如 'C-type' -> 'Fortran'), 而 flatten 返回的总是一个 copy. 现在你知道谁在拖你的后腿了吧! 下面的测试证明, 相比于 flatten, ravel 是神速.

def f1(a):
    for _ in range(N):
        a.flatten()

def f2(b):
    for _ in range(N):
        b.ravel()

print('%f' % ((t1-t0)/N))    # 0.001059
print('%f' % ((t2-t1)/N))    # 0.000000

选择数据

选择数据的时候, 我们常会用到 view 或者 copy 的形式. 我们知道了, 如果能用到 view 的, 我们就尽量用 view, 避免 copy 数据. 那什么时候会是 view 呢? 下面举例的都是 view 的方式:

a_view1 = a[1:2, 3:6]    # 切片 slice
a_view2 = a[:100]        # 同上
a_view3 = a[::2]         # 跳步
a_view4 = a.ravel()      # 上面提到了
...                      # 我只能想到这些, 如果还有请大家在评论里提出

那哪些操作我们又会变成 copy 呢?

a_copy1 = a[[1,4,6], [2,4,6]]   # 用 index 选
a_copy2 = a[[True, True], [False, True]]  # 用 mask
a_copy3 = a[[1,2], :]        # 虽然 1,2 的确连在一起了, 但是他们确实是 copy
a_copy4 = a[a[1,:] != 0, :]  # fancy indexing
a_copy5 = a[np.isnan(a), :]  # fancy indexing
...                          # 我只能想到这些, 如果还有请大家在评论里提出

Numpy 给了我们很多很自由的方式选择数据, 这些虽然都很方便, 但是如果你可以尽量避免这些操作, 你的速度可以飞起来.

在上面提到的 blog 里面, 他提到了, 如果你还是喜欢这种 fancy indexing 的形式, 我们也是可以对它加点速的. 那个 blog 中指出了两种方法

1. 使用 `np.take()`, 替代用 index 选数据的方法.

上面提到了如果用index 来选数据, 像 `a_copy1 = a[[1,4,6], [2,4,6]]`, 用 take 在大部分情况中会比这样的 `a_copy1` 要快.

a = np.random.rand(1000000, 10)
N = 99
indices = np.random.randint(0, 1000000, size=10000)

def f1(a):
    for _ in range(N):
        _ = np.take(a, indices, axis=0)

def f2(b):
    for _ in range(N):
        _ = b[indices]

print('%f' % ((t1-t0)/N))    # 0.000393
print('%f' % ((t2-t1)/N))    # 0.000569

2. 使用 `np.compress()`, 替代用 mask 选数据的方法.

上面的 `a_copy2 = a[[True, True], [False, True]]` 这种就是用 TRUE, FALSE 来选择数据的. 测试如下:

mask = a[:, 0] < 0.5
def f1(a):
    for _ in range(N):
        _ = np.compress(mask, a, axis=0)

def f2(b):
    for _ in range(N):
        _ = b[mask]

print('%f' % ((t1-t0)/N))    # 0.028109
print('%f' % ((t2-t1)/N))    # 0.031013

非常有用的 out 参数

不深入了解 numpy 的朋友, 应该会直接忽略很多功能中的这个 out 参数 (之前我从来没用过). 不过当我深入了解了以后, 发现他非常有用! 比如下面两个其实在功能上是没差的, 不过运算时间上有差, 我觉得可能是 `a=a+1` 要先转换成 `np.add()` 这种形式再运算, 所以前者要用更久一点的时间.

a = a + 1         # 0.035230
a = np.add(a, 1)  # 0.032738

如果是上面那样, 我们就会触发之前提到的 copy 原则, 这两个被赋值的 a, 都是原来 a 的一个 copy, 并不是 a 的 view. 但是在功能里面有一个 out 参数, 让我们不必要重新创建一个 a. 所以下面两个是一样的功能, 都不会创建另一个 copy. 不过可能是上面提到的那个原因, 这里的运算时间也有差.

a += 1                 # 0.011219
np.add(a, 1, out=a)    # 0.008843

带有 out 的 numpy 功能都在这里: Universal functions. 所以只要是已经存在了一个 placeholder (比如 a), 我们就没有必要去再创建一个, 用 out 方便又有效.

给数据一个名字

我喜欢用 pandas, 因为 pandas 能让你给数据命名, 用名字来做 index. 在数据类型很多的时候, 名字总是比 index 好记太多了, 也好用太多了. 但是 pandas 的确比 numpy 慢. 好在我们还是有途径可以实现用名字来索引. 这就是 structured array. 下面 a/b 的结构是一样的, 只是一个是 numpy 一个是 pandas.

a = np.zeros(3, dtype=[('foo', np.int32), ('bar', np.float16)])
b = pd.DataFrame(np.zeros((3, 2), dtype=np.int32), columns=['foo', 'bar'])
b['bar'] = b['bar'].astype(np.float16)

"""   
# a
array([(0,  0.), (0,  0.), (0,  0.)],
      dtype=[('foo', '<i4'), ('bar', '<f2')])

# b
   foo  bar
0    0  0.0
1    0  0.0
2    0  0.0
"""

def f1(a):
    for _ in range(N):
        a['bar'] *= a['foo']

def f2(b):
    for _ in range(N):
        b['bar'] *= b['foo']

print('%f' % ((t1-t0)/N))    # 0.000003
print('%f' % ((t2-t1)/N))    # 0.000508

可以看出来, numpy 明显比 pandas 快很多. 如果需要使用到不同数据形式, numpy 也是可以胜任的, 并且在还保持了快速的计算速度. 至于 pandas 为什么比 numpy 慢, 因为 pandas data 里面还有很多七七八八的数据, 记录着这个 data 的种种其他的特征. 这里还有更全面的对比: Numpy Vs Pandas Performance Comparison

如果大家还有其他的小技巧或者是速度大比拼, 欢迎在下面讨论. (一切为了速度~)

最后插个小广告, 如果你对机器学习感兴趣, 这里有很多厉害的短片形式机器学习方法介绍和很多机器学习的 Python 实践教程, 让你可以用业余时间秒懂机器学习: 莫烦Python 教程