NumPy 性能优化技巧

NumPy 以其高效的数组操作而闻名，但编写不当的代码仍可能造成性能瓶颈。掌握一些关键的优化技巧，可以充分发挥 NumPy 的速度优势，尤其是在处理大规模数据时。本章将介绍从基础到进阶的性能优化方法。

1. 向量化操作

向量化是 NumPy 性能的核心。尽可能使用 NumPy 内置的数组运算和通用函数（ufunc），避免显式的 Python 循环。

import numpy as np
import time

# 不推荐：Python 循环
def loop_add(a, b):
    result = np.empty_like(a)
    for i in range(len(a)):
        result[i] = a[i] + b[i]
    return result

# 推荐：向量化
def vec_add(a, b):
    return a + b

a = np.random.rand(1000000)
b = np.random.rand(1000000)

start = time.time()
loop_add(a, b)
print("Python 循环耗时:", time.time() - start)

start = time.time()
vec_add(a, b)
print("向量化耗时:", time.time() - start)

输出表明向量化通常快几十到上百倍。

2. 使用通用函数 (ufunc)

通用函数是对数组进行逐元素运算的 C 语言实现，速度极快。应优先使用 np.add、np.sin 等，而不是 Python 内置函数或自定义循环。

# 不推荐：使用 math.sin 循环
import math
result = [math.sin(x) for x in a]

# 推荐：使用 np.sin
result = np.sin(a)

3. 选择合适的数据类型

数据类型直接影响内存占用和计算速度。在满足精度要求的前提下，尽量使用较小的数据类型（如 float32 代替 float64，int8 代替 int64）。

arr_int64 = np.random.randint(0, 100, 1000000, dtype=np.int64)
arr_int8 = arr_int64.astype(np.int8)

# int8 占用的内存仅为 int64 的 1/8
print("int64 内存:", arr_int64.nbytes)
print("int8 内存:", arr_int8.nbytes)

4. 内存布局：C 风格与 Fortran 风格

NumPy 数组默认采用 C 风格（行主序）内存布局。对于某些操作（如沿行计算），C 风格更快；沿列计算则 Fortran 风格（列主序）可能更快。使用 order='F' 创建 Fortran 风格数组，或通过 np.asfortranarray() 转换。

# 创建 C 风格和 Fortran 风格的数组
c_arr = np.random.rand(1000, 1000)
f_arr = np.asfortranarray(c_arr)

# 沿行求和（C 风格更快）
%timeit c_arr.sum(axis=1)
%timeit f_arr.sum(axis=1)

# 沿列求和（Fortran 风格更快）
%timeit c_arr.sum(axis=0)
%timeit f_arr.sum(axis=0)

注：在交互式环境中可用 %timeit 测试。

5. 避免不必要的复制

尽量使用视图而非副本。切片、reshape、transpose 等通常返回视图，而花式索引、布尔索引返回副本。

arr = np.arange(1000000)

# 视图（不复制）
view = arr[::2]

# 副本（复制）
copy = arr[[0, 2, 4]]

# 检查是否共享内存
print("视图共享?", np.may_share_memory(arr, view))
print("副本共享?", np.may_share_memory(arr, copy))

6. 利用广播

广播允许不同形状的数组进行运算，避免了显式扩展数组（如 np.tile），从而节省内存和时间。

a = np.random.rand(1000, 100)
b = np.random.rand(100)

# 不推荐：手动扩展 b
b_expanded = np.tile(b, (1000, 1))
result = a + b_expanded

# 推荐：广播
result = a + b   # b 自动广播到每一行

7. 使用 out 参数减少内存分配

许多 NumPy 函数和 ufunc 支持 out 参数，允许将结果写入预先分配的数组，避免创建新数组。

a = np.random.rand(1000)
b = np.random.rand(1000)
result = np.empty_like(a)

# 使用 out 参数
np.add(a, b, out=result)   # 结果直接写入 result，无额外分配

8. 就地操作

使用 +=、*= 等运算符进行就地修改，避免创建新数组。

a = np.random.rand(1000)
b = np.random.rand(1000)

# 创建新数组
c = a + b

# 就地操作（修改 a）
a += b   # 更省内存

9. 合并操作以减少临时数组

将多个操作合并为一条表达式可以减少中间数组的产生。但也要注意可读性。

a = np.random.rand(1000)
b = np.random.rand(1000)
c = np.random.rand(1000)

# 多个中间数组
temp1 = a + b
temp2 = temp1 * c
result = np.sqrt(temp2)

# 合并操作
result = np.sqrt((a + b) * c)   # 更少的临时数组

10. 使用 numexpr 加速复杂表达式

对于复杂的算术表达式，numexpr 库可以并行计算并减少内存占用，尤其适合大型数组。

import numexpr as ne

a = np.random.rand(1000000)
b = np.random.rand(1000000)
c = np.random.rand(1000000)

# 使用 numexpr
result = ne.evaluate('sqrt((a + b) * c)')

11. 使用 numba JIT 编译

对于无法完全向量化的复杂算法，可以使用 numba 将 Python 函数即时编译为机器码，获得接近 C 的速度。

from numba import jit

@jit(nopython=True)
def my_func(a, b):
    result = np.empty_like(a)
    for i in range(len(a)):
        result[i] = a[i] ** 2 + b[i]   # 复杂逻辑
    return result

a = np.random.rand(1000000)
b = np.random.rand(1000000)
result = my_func(a, b)

12. 利用并行 BLAS 库

NumPy 的线性代数运算（如 np.dot、@）底层调用经过高度优化的 BLAS 库（如 OpenBLAS、MKL），它们会自动利用多核并行。确保你的 NumPy 链接了多线程 BLAS 以充分发挥多核性能。

13. 使用内存映射处理超大数据

当数据无法一次性装入内存时，使用 np.memmap 将磁盘文件映射为数组，按需加载。

# 创建内存映射文件（写入模式）
mmap = np.memmap('large.dat', dtype='float32', mode='w+', shape=(10000, 10000))
mmap[:, :] = 1.0   # 操作部分数据
mmap.flush()

# 只读模式读取
mmap_read = np.memmap('large.dat', dtype='float32', mode='r', shape=(10000, 10000))
print(mmap_read[0, :5])

14. 避免 np.vectorize 和 np.frompyfunc 的误用

虽然 np.vectorize 可以简化代码，但它本质是 Python 循环，不会带来性能提升。仅在无法向量化时作为便利手段使用。

# 不推荐用于性能
@np.vectorize
def slow_func(x):
    return x**2 if x > 0 else 0

# 应尽量使用向量化操作
def fast_func(x):
    return np.where(x > 0, x**2, 0)

15. 使用 np.einsum 优化张量运算

np.einsum 可以用 Einstein 求和约定表达复杂的张量收缩，有时比多个点积更高效。

A = np.random.rand(100, 100)
B = np.random.rand(100, 100)

# 矩阵乘法
C = np.einsum('ij,jk->ik', A, B)