使用Python玩转GPU

问题

随着机器学习对模型运算速度的需求越来越强烈，
一直想进行GPU编程，但一直以来这些都是c++的专利
一想到c++里的各种坑，就提不起劲来，毕竟这样来来回回填坑的投入产出，生产效率就会大打折扣

解决方案

让人欣喜的是，随着Python阵营的不断发展壮大，使用python进行GPU编程也越来越便捷了

那么具体有些什么样的包，能针对GPU做些啥事呢？
看看一些具体的代码，就能大概明白：
首先是pycuda，这是它的一个例子：

mod = SourceModule(""" __global__ void multiply_them(float *dest, float *a, float *b) { const int i = threadIdx.x; dest[i] = a[i] * b[i]; } """)

由上面的代码我们可以看出，pycuda将调用gpu的c++代码做了包装，可以在python里直接使用

再看看numba：

@cuda.jit def increment_by_one(an_array): pos = cuda.grid(1) if pos < an_array.size: an_array[pos] += 1

我们可以发现，numba更进一步，直接使用装饰器的办法让调用GPU的过程更简洁方便

再看看cupy：

import numpy as np import cupy as cp x_gpu = cp.array([1, 2, 3]) l2_gpu = cp.linalg.norm(x_gpu)

cupy的调用方法看起来更加简单清晰，直接将np替换为cp即可

比较

所以，从机器学习全流程的角度我做了下汇总：

原始框架	GPU替代包	支持GPU	C/C++核函数	表达式核函数	装饰器 & Python
cuda	pycuda	n卡	√	×	×
opencl	pyopencl	n卡 + a卡	√	×	×
numpy	numba	n卡	√	√	√
scipy	cupy	n卡	√	√	√
pandas	cudf & modin	n卡	无(自动)
sklearn	cuml & scikit-cuda	n卡	无(自动)

目前cupy和numba对numpy的支持都不全面，可两者结合使用

从上面可以看出，基本上已经涵盖了机器学习的全流程，大部分包只支持cuda，主要都是为方便使用n卡加速
部分包还是只能使用c/c++语言构建核函数，主要还是受限于cuda驱动的capability
所以实际使用门槛并没有降低，只是将核函数包装到python里使用

GPU的主要优势在于大规模的并行计算，所以我又收集了一些并行计算框架，方便日后使用

框架	CPU并行计算	分布式并行计算	GPU并行计算
Multiprocess	√	×	×
joblib	√	×	×
dask	√	√	√
ray	√	√	√

至于机器学习/深度学习框架，那就更不用说了：
xgboost已经开始支持GPU，RandomForest也有GPU的版本，
tensorflow, pytorch默认就是支持GPU的，这里就不再赘述

结论

没有框架时，我们希望有效率的python包能快速解决问题，可框架多了，又会出现选择困难症
各种框架都宣称自己效率高，灵活好用，各种benchmark也让人眼花缭乱，目不暇接
到底用哪个框架合适，我把自己的一些经验也总结下，希望能让大家少踩一些坑：

1.对于一般的并行计算任务，使用joblib就能方便完成；

2.对于需要集群或GPU的计算任务，可以选择dask或ray；
这里推荐下dask，dask于机器学习/深度学习的计算包结合的更紧密，推出了dask_ml用于处理分布式机器学习；

3.如果想快速迁移numpy/pandas的代码到gpu，可以使用cupy + cudf的组合方式；

4.如果有复杂自定义的计算以及为了追求性能，可以使用pycuda + numba的形式；

5.对于numpy的替换到底选用cupy还是numba？
这里没有严格的界限，两者对GPU的调用方式设计，实际都会有一定的编码成本
从cupy的基本例子中可以看出，对于部分调用来说cupy更简洁，但是牺牲了cpu并行和分布式并行的功能为代价
所以目前可以持续关注这两个框架

6.对于替换pandas到底选用cudf还是modin？
modin本身并不是专为cuda并行化而设计，它只是底层支持了dask和ray，由此间接的支持了GPU
且到目前为止对pandas方法的支持还不全面，所以这里推荐选择cudf

思考

总之，python作为机器学习的首选语言，正在不断的开疆拓土，不断的降低并行计算的门槛
短短几年前，还只能用xgboost + spark的方式进行分布式训练，转眼现在就有了多种python解决方案
短短几年前，还只能用c++ cuda的方式进行GPU编程，转眼现在也有了多种纯python的框架支持
短短几年前，将GPU进行集群化、虚拟化管理几乎是不可能的，转眼现在也有了可靠的解决方案
……

但是，到目前为止，还没有一款真正能充分智能化的利用并行能力计算的框架：它能综合cpu+gpu+分布式的计算能力，目标就是为了加速计算，得到结果。期待这样的框架诞生！

相信不久的将来，会有更多更强大的python框架出现，不断的加速自动化的进程
让更多的生产力能从原始的轮子中解放出来，加快人工智能的进化！