Chuoer47's Blog

Appearance

性能与技巧 ​

NumPy 的高性能来源于向量化运算和高效的内存管理。本章介绍如何编写高效的 NumPy 代码，以及常见的性能优化技巧。

一、向量化思维 ​

1、避免 Python 循环 ​

向量化是 NumPy 性能的核心。将显式循环替换为数组操作，通常能获得 10~100 倍的速度提升。

python
import numpy as np

# ❌ 慢：Python 循环
def slow_distance(points):
    n = len(points)
    dist = np.zeros((n, n))
    for i in range(n):
        for j in range(n):
            dist[i, j] = np.sqrt(np.sum((points[i] - points[j]) ** 2))
    return dist

# ✅ 快：向量化
def fast_distance(points):
    diff = points[:, np.newaxis, :] - points[np.newaxis, :, :]
    return np.sqrt(np.sum(diff ** 2, axis=2))

points = np.random.randn(1000, 2)
# fast_distance 比 slow_distance 快约 100 倍

2、向量化替代条件逻辑 ​

python
a = np.array([1, -2, 3, -4, 5])

# ❌ 慢
result = np.zeros_like(a)
for i in range(len(a)):
    if a[i] > 0:
        result[i] = a[i] ** 2
    else:
        result[i] = 0

# ✅ 快：np.where
result = np.where(a > 0, a ** 2, 0)

# ✅ 更复杂的条件：np.select
conditions = [a < -2, a < 0, a < 3, a >= 3]
choices     = [-100,    -10,   100,   1000]
result = np.select(conditions, choices, default=0)

3、向量化替代搜索 ​

python
data = np.random.randn(10000)

# ❌ 慢：找满足条件的索引
indices = []
for i in range(len(data)):
    if data[i] > 2.0:
        indices.append(i)

# ✅ 快
indices = np.where(data > 2.0)[0]
# 或
indices = np.nonzero(data > 2.0)[0]

二、内存管理 ​

1、视图 vs 副本 ​

python
a = np.arange(1000000)

# ❌ 创建副本（额外内存）
b = a.copy()

# ✅ 使用视图（共享内存，零开销）
b = a.view()
c = a[1:100:2]      # 切片是视图
d = a.reshape(1000, 1000)  # reshape 可能返回视图

# 检查是否共享内存
print(np.shares_memory(a, b))  # True

2、预分配内存 ​

python
# ❌ 慢：循环中拼接
result = np.array([])
for i in range(10000):
    result = np.append(result, i * 2)  # 每次创建新数组！

# ✅ 快：预分配
result = np.empty(10000)
for i in range(10000):
    result[i] = i * 2

# ✅✅ 最快：向量化
result = np.arange(10000) * 2

3、选择合适的数据类型 ​

python
# 默认 float64（8 字节/元素）
a = np.zeros(1000000)                  # 8 MB

# 如果精度够用，float32 省一半内存
b = np.zeros(1000000, dtype=np.float32)  # 4 MB

# 整数同理
c = np.zeros(1000000, dtype=np.int8)     # 1 MB

# 比较
print(a.nbytes)  # 8000000
print(b.nbytes)  # 4000000
print(c.nbytes)  # 1000000

精度 vs 性能权衡

深度学习中常用 float32 甚至 float16，在精度可接受的范围内大幅减少内存占用和计算时间。只有在需要高精度计算（如求解线性方程组）时才使用 float64。

4、内存布局优化 ​

python
# C-order（行优先，默认）：利于按行操作
a = np.zeros((1000, 1000), order='C')

# Fortran-order（列优先）：利于按列操作
b = np.zeros((1000, 1000), order='F')

# 按行求和 → C-order 更快
# 按列求和 → F-order 更快

# 检查布局
print(a.flags.c_contiguous)  # True
print(b.flags.f_contiguous)  # True

# 转换布局
c = np.ascontiguousarray(b)  # 转为 C-order
d = np.asfortranarray(a)     # 转为 F-order

三、ufunc 高级技巧 ​

1、out 参数（避免临时数组） ​

python
a = np.arange(1000000, dtype=np.float64)
b = np.arange(1000000, dtype=np.float64)

# ❌ 创建临时数组
c = a + b

# ✅ 直接写入预分配的输出数组
c = np.empty_like(a)
np.add(a, b, out=c)

2、reduce 与 accumulate ​

python
a = np.arange(1, 6)  # [1, 2, 3, 4, 5]

# reduce: 沿轴累积运算（返回标量或降维数组）
print(np.add.reduce(a))        # 15  (等价于 np.sum)
print(np.multiply.reduce(a))   # 120 (等价于 np.prod)

# accumulate: 返回中间结果
print(np.add.accumulate(a))        # [ 1  3  6 10 15]
print(np.multiply.accumulate(a))   # [  1   2   6  24 120]

3、outer（外积） ​

python
a = np.array([1, 2, 3])
b = np.array([10, 20])

# 外积
print(np.multiply.outer(a, b))
# [[10 20]
#  [20 40]
#  [30 60]]

# 乘法表
x = np.arange(1, 10)
print(np.multiply.outer(x, x))

4、at（高级索引赋值） ​

python
a = np.zeros(10)
indices = [1, 3, 5, 3, 1]
values = [10, 20, 30, 40, 50]

# ❌ 简单索引赋值：重复索引只保留最后一次
a[indices] = values
print(a[1])   # 50（不是 10+50=60）

# ✅ np.add.at: 对重复索引累加
a = np.zeros(10)
np.add.at(a, indices, values)
print(a[1])   # 60 (10 + 50)
print(a[3])   # 60 (20 + 40)

四、高效操作模式 ​

1、einsum（爱因斯坦求和） ​

np.einsum 是一个强大的函数，可以用简洁的下标表示法表达复杂的多维运算：

python
A = np.random.randn(3, 4)
B = np.random.randn(4, 5)

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

# 转置
D = np.einsum('ij->ji', A)          # 等价于 A.T

# 迹（对角线之和）
E = np.random.randn(3, 3)
trace = np.einsum('ii->', E)        # 等价于 np.trace(E)

# 批量矩阵乘法
batch_A = np.random.randn(100, 3, 4)
batch_B = np.random.randn(100, 4, 5)
batch_C = np.einsum('bij,bjk->bik', batch_A, batch_B)

# 向量点积
x = np.array([1, 2, 3])
y = np.array([4, 5, 6])
dot = np.einsum('i,i->', x, y)     # 等价于 np.dot(x, y)

# 外积
outer = np.einsum('i,j->ij', x, y)

2、np.einsum 索引规则 ​

表达式含义等价操作
'ij,jk->ik'矩阵乘法A @ B
'ij->ji'转置A.T
'ii->'迹np.trace(A)
'ij->'全部求和np.sum(A)
'ij->i'按行求和np.sum(A, axis=1)
'ij->j'按列求和np.sum(A, axis=0)
'i,i->'点积np.dot(a, b)
'i,j->ij'外积np.outer(a, b)
'ij,ij->ij'逐元素乘A * B
'ij,ij->'全部乘积和np.sum(A * B)

einsum 何时使用

当一个运算需要多步中间操作时（如 np.sum(A * B.T, axis=1)），einsum 通常更高效且更可读。但简单运算（如 A + B）不需要 einsum。

3、分块计算（处理超大数组） ​

python
# 当数组太大无法一次装入内存时
def chunked_operation(filename, chunk_size=10000):
    """逐块读取并处理大文件"""
    total = 0
    count = 0
    for chunk in np.loadtxt(filename, delimiter=',',
                            chunksize=chunk_size):
        total += np.sum(chunk)
        count += chunk.size
    return total / count

# 对于内存中的大数组
large = np.random.randn(10000000)
result = np.empty_like(large)
chunk = 100000
for i in range(0, len(large), chunk):
    end = min(i + chunk, len(large))
    result[i:end] = np.sin(large[i:end]) ** 2 + np.cos(large[i:end]) ** 2

五、性能分析工具 ​

1、计时 ​

python
import time

a = np.random.randn(1000, 1000)
b = np.random.randn(1000, 1000)

# 方法一：time.time
start = time.time()
c = a @ b
elapsed = time.time() - start
print(f"耗时: {elapsed:.4f} 秒")

# 方法二：IPython / Jupyter 中的 %timeit
# %timeit a @ b

# 方法三：timeit 模块
import timeit
elapsed = timeit.timeit(lambda: a @ b, number=10)
print(f"10 次平均: {elapsed/10:.4f} 秒")

2、内存分析 ​

python
# 查看数组内存占用
a = np.random.randn(1000, 1000)
print(f"内存: {a.nbytes / 1024 / 1024:.2f} MB")  # 7.63 MB

# 检查是否是视图
print(a.base is None)     # True  — 是独立数组
b = a[::2, ::2]
print(b.base is None)     # False — 是视图

# 共享内存检测
print(np.shares_memory(a, b))  # True

3、性能优化清单 ​

优化策略适用场景
向量化替代所有 Python 循环
预分配 np.empty已知输出大小时
out 参数避免 ufunc 创建临时数组
视图替代副本切片、reshape、转置
合适的 dtypefloat32 替代 float64
内存布局行操作用 C-order，列操作用 F-order
einsum复杂多维运算
np.where替代条件循环
np.unique + return_index去重同时获取索引
批量操作一次处理一批而非逐个

六、常见陷阱 ​

1、可变对象作为默认参数 ​

python
# ❌ 危险：默认参数是可变对象
def bad_func(arr=[]):
    arr.append(1)
    return arr

# ✅ 安全
def good_func(arr=None):
    if arr is None:
        arr = []
    arr.append(1)
    return arr

2、整数溢出 ​

python
# NumPy 整数有固定范围，不像 Python int 自动扩展
a = np.array([200], dtype=np.uint8)
b = a + 100
print(b)        # [44]  — 溢出！200 + 100 = 300 % 256 = 44

# 解决：使用更大的类型
c = a.astype(np.int32) + 100
print(c)        # [300]

3、浮点精度 ​

python
# 浮点数比较不要用 ==
a = np.array([0.1 + 0.2])
b = np.array([0.3])
print(a == b)                # [False]
print(np.isclose(a, b))      # [True]
print(np.allclose(a, b))     # True

4、数组拷贝陷阱 ​

python
a = np.array([1, 2, 3])
b = a           # 不是拷贝！b 和 a 指向同一对象
b[0] = 99
print(a)        # [99  2  3]  — a 也被修改了

# 正确拷贝
c = a.copy()
c[0] = 0
print(a)        # [99  2  3]  — a 不受影响

📖 官方文档