7 个 NumPy 技巧助您实现代码向量化

📢 转载信息

原文链接：https://machinelearningmastery.com/7-numpy-tricks-to-vectorize-your-code/

原文作者：Bala Priya C

---

您可能编写了处理数据循环的 Python 代码。它清晰、正确，但在处理真实世界数据量时却慢得无法使用。问题不在于您的算法；而在于 Python 的 for 循环以解释器速度执行，这意味着每次迭代都要支付 Python 动态类型检查和内存管理的开销。

NumPy 有助于解决这个瓶颈。它封装了高度优化的 C 和 Fortran 库，这些库能够以单个操作处理整个数组，完全绕过了 Python 的开销。但您需要以不同的方式编写代码——并将其表达为向量化操作——才能获得这种速度。这种转变需要一种不同的思维方式。您不再是思考“遍历并检查每个值”，而是思考“选择匹配某个条件的元素”。您不再使用嵌套迭代，而是思考数组维度和广播机制。

本文将介绍 7 种消除数值代码中循环的向量化技术。每种技术都针对开发人员通常会求助于迭代的特定模式，向您展示如何用数组操作来重新构建问题。结果是代码运行速度快得多（非常快），而且通常比基于循环的版本更清晰易读。

🔗 GitHub 上的代码链接

1. 使用布尔索引代替条件循环

您需要根据条件过滤或修改数组元素。本能反应是循环并检查每一个元素。

import numpy as np # 慢：基于循环的过滤
data = np.random.randn(1000000)
result = []
for x in data:
    if x > 0:
        result.append(x * 2)
    else:
        result.append(x)
result = np.array(result)

这是向量化方法：

# 快：布尔索引
data = np.random.randn(1000000)
result = data.copy()
result[data > 0] *= 2

在这里，data > 0 创建了一个布尔数组——条件成立的位置为 True，否则为 False。将其用作索引，只选择满足条件的那些元素。

2. 使用广播进行隐式循环

有时您希望组合不同形状的数组，比如将一个行向量加到矩阵的每一行。基于循环的方法需要显式迭代。

# 慢：显式循环
matrix = np.random.rand(1000, 500)
row_means = np.mean(matrix, axis=1)
centered = np.zeros_like(matrix)
for i in range(matrix.shape[0]):
    centered[i] = matrix[i] - row_means[i]

这是向量化方法：

# 快：广播
matrix = np.random.rand(1000, 500)
row_means = np.mean(matrix, axis=1, keepdims=True)
centered = matrix - row_means

在此代码中，设置 keepdims=True 会使 row_means 的形状保持为 (1000, 1)，而不是 (1000,)。当您进行减法操作时，NumPy 会自动将此列向量拉伸到矩阵的所有列上。形状不匹配，但 NumPy 通过沿单例维度重复值来使它们兼容。

🔖 注意：当维度兼容时（即相等，或其中一个为 1），广播机制才会生效。较小的数组会虚拟地重复以匹配较大数组的形状，无需内存复制。

3. 使用 np.where() 实现向量化的 If-Else

当您需要根据条件对不同元素执行不同计算时，需要在循环中编写分支逻辑。

# 慢：循环中的条件逻辑
temps = np.random.uniform(-10, 40, 100000)
classifications = []
for t in temps:
    if t < 0:
        classifications.append('freezing')
    elif t < 20:
        classifications.append('cool')
    else:
        classifications.append('warm')

这是向量化方法：

# 快：np.where() 和 np.select()
temps = np.random.uniform(-10, 40, 100000)
classifications = np.select(
    [temps < 0, temps < 20, temps >= 20],
    ['freezing', 'cool', 'warm'],
    default='unknown' # 添加了字符串默认值
)

# 对于简单的分割，np.where() 更简洁：
scores = np.random.randint(0, 100, 10000)
results = np.where(scores >= 60, 'pass', 'fail')

np.where(condition, x, y) 在条件为 True 时返回 x 中的元素，否则返回 y 中的元素。np.select() 将此扩展到多个条件。它按顺序检查每个条件，并返回第二个列表中对应的值。

🔖 注意：np.select() 中的条件应是互斥的。如果一个元素的多个条件都为 True，则第一个匹配的条件获胜。

4. 更好的索引用于查找操作

假设您有一组索引，需要从多个位置收集元素。您通常会倾向于在循环中使用字典查找，或者更糟的是，进行嵌套搜索。

# 慢：基于循环的收集
lookup_table = np.array([10, 20, 30, 40, 50])
indices = np.random.randint(0, 5, 100000)
results = []
for idx in indices:
    results.append(lookup_table[idx])
results = np.array(results)

这是向量化方法：

lookup_table = np.array([10, 20, 30, 40, 50])
indices = np.random.randint(0, 5, 100000)
results = lookup_table[indices]

当您使用另一个整数数组对数组进行索引时，NumPy 会提取那些位置的元素。这在多个维度中也适用：

matrix = np.arange(20).reshape(4, 5)
row_indices = np.array([0, 2, 3])
col_indices = np.array([1, 3, 4])
values = matrix[row_indices, col_indices] # 获取 matrix[0,1], matrix[2,3], matrix[3,4]

🔖 注意：这在实现分类编码、构建直方图或任何将索引映射到值的操作中特别有用。

5. 使用 np.vectorize() 处理自定义函数

您有一个作用于标量的函数，但需要将其应用于数组。在各处编写循环会使代码变得冗长。

# 慢：手动循环
def complex_transform(x):
    if x < 0:
        return np.sqrt(abs(x)) * -1
    else:
        return x ** 2
data = np.random.randn(10000)
results = np.array([complex_transform(x) for x in data])

这是向量化方法：

# 更简洁：np.vectorize()
def complex_transform(x):
    if x < 0:
        return np.sqrt(abs(x)) * -1
    else:
        return x ** 2

vec_transform = np.vectorize(complex_transform)
data = np.random.randn(10000)
results = vec_transform(data)

在这里，np.vectorize() 包装了您的函数，使其能够处理数组。它会自动逐元素应用函数并处理输出数组的创建。

🔖 注意：这并不会神奇地让您的函数运行得更快。在底层，它仍然在 Python 中进行循环。这里的优势在于代码的清晰度，而非速度。要实现真正的性能提升，请直接使用 NumPy 操作重写函数：

# 真正快
data = np.random.randn(10000)
results = np.where(data < 0, -np.sqrt(np.abs(data)), data ** 2)

6. 使用 np.einsum() 处理复杂的数组操作

矩阵乘法、转置、迹和张量收缩等操作堆叠起来会形成一长串难以阅读的操作链。

# 标准方式的矩阵乘法
A = np.random.rand(100, 50)
B = np.random.rand(50, 80)
C = np.dot(A, B)

# 批处理矩阵乘法 - 会变得很乱
batch_A = np.random.rand(32, 10, 20)
batch_B = np.random.rand(32, 20, 15)
results = np.zeros((32, 10, 15))
for i in range(32):
    results[i] = np.dot(batch_A[i], batch_B[i])

这是向量化方法：

# 使用 np.einsum 进行批处理矩阵乘法
# 'b' 是批次维度，'i,j' 是矩阵维度
batch_C = np.einsum('bij,bjk->bik', batch_A, batch_B)

np.einsum()（爱因斯坦求和约定）提供了一种简洁的语法来表达复杂的张量操作，它清楚地表明了哪些维度应该被收缩（相乘并求和）以及结果维度应该是什么。

7. 使用 np.apply_along_axis() 处理更复杂的操作

当您的操作必须跨越特定轴（例如，找到每一行的最小值），但无法使用标准广播或元素级操作时，可以使用 np.apply_along_axis()。

# 慢：逐行查找最大值
data = np.random.rand(1000, 100)
max_values = []
for row in data:
    max_values.append(np.max(row))
max_values = np.array(max_values)

这是向量化方法：

# 快：使用 np.apply_along_axis
data = np.random.rand(1000, 100)
max_values = np.apply_along_axis(np.max, axis=1, arr=data)

np.apply_along_axis(func, axis, arr) 对数组的指定轴的每个一维切片应用函数。虽然这本质上是循环，但它在 C 级别上比显式 Python 循环快得多。

总结

向量化是将 Python 代码转换为高性能科学计算的关键一步。通过将思维从逐个元素迭代转向数组级操作，您可以显著减少代码执行时间。

以下是本文介绍的 7 个技巧的快速回顾：

• 布尔索引：用于基于条件的元素选择和修改，替代 if/else 循环。
• 广播：自动对齐不同形状的数组，避免显式循环。
• np.where() / np.select()：用于向量化的条件赋值。
• 高级索引：使用索引数组一步获取多个元素，替代查找循环。
• np.vectorize()：使标量函数能够处理数组，提高代码清晰度（但并非性能提升）。
• np.einsum()：用于简洁地表达复杂的张量和矩阵运算。
• np.apply_along_axis()：在需要跨轴应用复杂函数时使用。

---

🚀 想要体验更好更全面的AI调用？

欢迎使用青云聚合API，约为官网价格的十分之一，支持300+全球最新模型，以及全球各种生图生视频模型，无需翻墙高速稳定，文档丰富，小白也可以简单操作。