如何加速Python的数值计算?

147 阅读5分钟

数据准备

首先,我们需要创建一个包含浮点数的文件,格式可以是简单的文本格式,例如 CSV。内容可以是随机的一维浮点数组:

我们可以使用以下代码生成一个包含浮点数的文件(假设命名为 data.txt):

import numpy as np

size = 10000000
data_a = np.random.rand(size)
data_b = np.random.rand(size)

# 将数据写入文件
np.savetxt('data_a.txt', data_a)
np.savetxt('data_b.txt', data_b)

基础版本对比

Python 实现

接下来是 Python 版本的实现,分阶段统计耗时:

import numpy as np
import time


def read_data(filename):
    return np.loadtxt(filename)


def cosine_similarity(a, b):
    dot_product = np.dot(a, b)  # 使用 NumPy 的 dot 函数计算点积
    norm_a = np.linalg.norm(a)
    norm_b = np.linalg.norm(b)
    
    # 余弦相似度
    similarity = dot_product / (norm_a * norm_b)
    return similarity


# 读数据
start_read = time.time()
a = read_data('data_a.txt')
b = read_data('data_b.txt')
end_read = time.time()

# 计算相似度
start_calc = time.time()
similarity = cosine_similarity(a, b)
end_calc = time.time()

print(f"Python Cosine Similarity: {similarity}")
print(f"Python Reading Time: {end_read - start_read:.6f} seconds")
print(f"Python Calculation Time: {end_calc - start_calc:.6f} seconds")

Python Cosine Similarity: 0.7500025214194538
Python Reading Time: 16.402409 seconds
Python Calculation Time: 0.039559 seconds

Python版本优化

优化后的 Python 代码示例

1. 使用 NumPy 的向量化运算

NumPy 本身就是为了处理数组和矩阵计算而优化的,完整利用 NumPy 的向量化运算可以显著提高性能。我们已经在前面的代码示例中使用了 NumPy,但可以更进一步,例如,通过避免显式的循环。

2. 避免重复计算

在计算余弦相似度时,我们可以避免重复计算向量的范数。缓存它们会更高效。

3. 使用 @ 操作符

在计算向量点积时,使用 @ 操作符(矩阵乘法)会更容易并且通常更快。

以下是经过优化后的代码:

import numpy as np
import time

def read_data(filename):
    return np.loadtxt(filename)

# 读数据
start_read = time.time()
a = read_data('data_a.txt')
b = read_data('data_b.txt')
end_read = time.time()

# 计算相似度
start_calc = time.time()
# 计算点积和范数
dot_product = np.dot(a, b)
norm_a = np.linalg.norm(a)
norm_b = np.linalg.norm(b)

# 余弦相似度
similarity = dot_product / (norm_a * norm_b)
end_calc = time.time()

print(f"Python Cosine Similarity: {similarity}")
print(f"Python Reading Time: {end_read - start_read:.6f} seconds")
print(f"Python Calculation Time: {end_calc - start_calc:.6f} seconds")

使用多进程计算

import numpy as np
import time
from concurrent.futures import ProcessPoolExecutor


def cosine_similarity(pair):
    a, b = pair
    dot_product = np.dot(a, b)
    norm_a = np.linalg.norm(a)
    norm_b = np.linalg.norm(b)
    return dot_product / (norm_a * norm_b)

def read_data(filename):
    return np.loadtxt(filename)


if __name__ == "__main__":
    start_read = time.time()
    a = read_data('data_a.txt')
    b = read_data('data_b.txt')
    end_read = time.time()

    # 将向量配对
    vector_pairs = [(a, b)]  # 或者多个向量可以组合成对

    start_calc = time.time()
    with ProcessPoolExecutor(max_workers=4) as pool:  # 使用4个进程
        similarities = pool.map(cosine_similarity, vector_pairs)
    end_calc = time.time()

    print(f"Similarities: {similarities}")
    print(f"Reading Time: {end_read - start_read:.6f} seconds")
    print(f"Calculation Time: {end_calc - start_calc:.6f} seconds")

这里会设计到多进程的管理开销,并不能解决单个函数计算的效率。

使用Cython对计算进行加速

使用 Cython 进行优化通常是通过使用 C 接口、Cython 或其他扩展模块来实现更高效的代码执行。在 Python 的上下文中,Cython 是默认的实现,并且它的性能瓶颈主要来自于几个方面,如 GIL(全局解释器锁)、解释执行和内存管理等。

pip install cython

编写 Cython 模块

创建一个 cosine_similarity.pyx 文件,内容如下:

import numpy as np
cimport numpy as cnp
from libc.math cimport sqrt

cpdef double cosine_similarity(cnp.ndarray[cnp.double_t, ndim=1] a, 
                               cnp.ndarray[cnp.double_t, ndim=1] b):
    cdef double dot_product = 0.0
    cdef double norm_a = 0.0
    cdef double norm_b = 0.0
    cdef Py_ssize_t i

    for i in range(a.shape[0]):
        dot_product += a[i] * b[i]
        norm_a += a[i] * a[i]
        norm_b += b[i] * b[i]

    return dot_product / (sqrt(norm_a) * sqrt(norm_b))

编写 setup.py

创建一个 setup.py 文件,用于编译 Cython 代码:

from setuptools import setup
from Cython.Build import cythonize
import numpy as np

setup(
    ext_modules=cythonize("cosine_similarity.pyx"),
    include_dirs=[np.get_include()]
)

编译 Cython 模块

python setup.py build_ext --inplace

使用 Cython 的优化版本

import numpy as np
import time
from cosine_similarity import cosine_similarity  # 导入 Cython 函数


def read_data(filename):
    return np.loadtxt(filename)


start_read = time.time()
a = read_data('data_a.txt')
b = read_data('data_b.txt')
end_read = time.time()

start_calc = time.time()
similarities = cosine_similarity(a, b)
end_calc = time.time()

print(f"Similarities: {similarities}")
print(f"Reading Time: {end_read - start_read:.6f} seconds")
print(f"Calculation Time: {end_calc - start_calc:.6f} seconds")


Similarities: 0.7500025214194899
Reading Time: 14.760010 seconds
Calculation Time: 0.028212 seconds

使用Numba对计算进行加速

Numba 是一个高性能的 JIT 编译器,可以快速将 Python 代码编译为机器代码。它在计算密集型计算时通常表现良好。

pip install numba

使用 Numba 将计算函数标记为 @njit(即时编译):

import numpy as np
import time
from numba import njit

@njit
def cosine_similarity(a, b):
    dot_product = 0.0
    norm_a = 0.0
    norm_b = 0.0
    for i in range(len(a)):
        dot_product += a[i] * b[i]
        norm_a += a[i] * a[i]
        norm_b += b[i] * b[i]
    return dot_product / (np.sqrt(norm_a) * np.sqrt(norm_b))

def read_data(filename):
    return np.loadtxt(filename)


if __name__ == "__main__":
    start_read = time.time()
    a = read_data('data_a.txt')
    b = read_data('data_b.txt')
    end_read = time.time()

    start_calc = time.time()
    similarity = cosine_similarity(a, b)
    end_calc = time.time()

    print(f"Python Cosine Similarity: {similarity}")
    print(f"Reading Time: {end_read - start_read:.6f} seconds")
    print(f"Calculation Time: {end_calc - start_calc:.6f} seconds")

虽然使用了Numba进行编译加速,但是由于使用了for 训练来进行计算,因此效率并不高。

import time
import numpy as np
from numba import njit

def read_data(filename):
    return np.loadtxt(filename)

@njit
def cosine_similarity(a, b):
    dot_product = np.dot(a, b)  # 使用 NumPy 的 dot 函数计算点积
    norm_a = np.linalg.norm(a)
    norm_b = np.linalg.norm(b)
    
    # 余弦相似度
    similarity = dot_product / (norm_a * norm_b)
    return similarity


start_read = time.time()
a = read_data('data_a.txt')
b = read_data('data_b.txt')
end_read = time.time()

start_calc = time.time()
similarity = cosine_similarity(a, b)
end_calc = time.time()

print(f"Python Cosine Similarity: {similarity}")
print(f"Reading Time: {end_read - start_read:.6f} seconds")
print(f"Calculation Time: {end_calc - start_calc:.6f} seconds")

使用 @njit 的装饰器来加速 NumPy 计算时,通常期望通过 JIT(即时编译)技术来提高性能。然而,使用 Numba 的性能提升并不是绝对的,可能会有以下一些原因导致运行速度变慢:

  1. 小数据规模

对于小规模的数据(如长度为几百或几千的数组),JIT 编译所需的初始编译时间可能超过了计算本身的时间,从而导致整体运行时间变长。这是因为 JIT 编译会在首次调用时进行编译,之后的调用才会受益于加速。

  1. NumPy 限制

NumPy 的一部分功能(如某些数据类型和处理)可能无法高效地在 Numba 中使用。如果使用的 NumPy 函数不被 Numba 支持,它们将无法利用 JIT 加速。

  1. 函数重排与并行

如果你在 JIT 编译中使用了较复杂的逻辑,Numba 的优化可能不如直接使用 NumPy 的矢量化操作效果好。

  1. 内存管理

在某些情况下,NumPy 的内存管理可能更高效。NumPy 使用 C 后端进行数组操作,这些操作经过高度优化。

在选择用 NumPy 还是用 Numba 进行优化时,最好先测量一下,看看哪种方法在特定情况下表现更好。如果你的数据量大,NumPy 和 Numba 的性能差异会越来越明显;但对于小数据集,可能 NumPy 已经足够快。