用GPU加速python代码的运算速度

用GPU加速python代码的运算速度

1. Numba是什么？

Numba是一个库，可以在运行时将Python代码编译为本地机器指令，而不会强制大幅度的改变普通的Python代码（稍后再做说明）。翻译/魔术是使用LLVM编译器完成的，该编译器是相当活跃的开源社区开发的。

Numba最初是由Continuum Analytics内部开发，此公司也开发了著名的Anaconda，但现在它是开源的。核心应用领域是math-heavy（密集数学？重型数学？）和array-oriented（面向数组）的功能，它们在本地Python中相当缓慢。想象一下，在Python中编写一个模块，必须一个元素接着一个元素的循环遍历一个非常大的数组来执行一些计算，而不能使用向量操作来重写。这是很不好的主意，是吧？所以“通常”这类库函数是用C / C ++或Fortran编写的，编译后，在Python中作为外部库使用。Numba这类函数也可以写在普通的Python模块中，而且运行速度的差别正在逐渐缩小。

2.怎么才能get到Numba呢？

安装Numba的推荐方法是使用conda包管理

conda install numba

3.如何使用Numba呢？

使用它的要求不多。基本上，你写一个自己的“普通”的Python函数，然后给函数定义添加一个装饰（如果你不是很熟悉装饰器，读一下关于this或that）。你可以使用不同类型的装饰器，但@jit可能是刚开始的选择之一。其他装饰器可用于例如创建numpy通用功能@vectorize或编写将在CUDA GPU上执行的代码@cuda。
具体使用方法：from numba import jit导入包后，在方法前加标注：@jit即可

代码案例：

import numpy as np
from numba import jit
import time

@jit
def sum_jit(arr):
    s_time = time.time()
    m = arr.shape[0]
    result = 0.0
    for i in range(m):
        result += arr[i]
    e_time = time.time()
    return (e_time-s_time)

def sum(arr):
    s_time = time.time()
    m = arr.shape[0]
    result = 0.0
    for i in range(m):
        result += arr[i]
    e_time = time.time()
    return (e_time-s_time)

def main():
    n = int(10.0*1e6)
    array = np.random.random(n)
    t1 = sum_jit(array)
    t2 = sum(array)
    print("Time with JIT:", t1)
    print("Time without JIT:", t2)

if __name__ == '__main__':
    main()