Jetson學習筆記(七):使用GPU加速代碼

首先，在Jetson上如何安裝Numba？相信大家現在刷的JetPack應該都是4.X版本，基于Ubuntu 18.04，那麼直接用apt安裝就可以了。

sudo apt-get install python3-numba
           

這樣Numba就已經裝好了，不過此時隻能生成CPU上的代碼。為了讓Numba能正确找到GPU以及所需的各種CUDA庫，我們還需要配置一些環境變量。我推薦直接使用超級使用者權限修改/etc/environment，全局配置，一勞永逸。在/etc/environment中添加如下環境變量配置。

NUMBAPRO_CUDALIB="/usr/local/cuda/lib64"
NUMBAPRO_LIBDEVICE="/usr/local/cuda/nvvm/libdevice"
NUMBAPRO_NVVM="/usr/local/cuda/nvvm/lib64/libnvvm.so"
           

如果要使用Jupyter，在Service檔案中也配置一下EnvironmentFile字段就行了。

EnvironmentFile=/etc/environment
           

重新開機後應該就生效了，可以在終端執行下面的指令确認一下。

numba -s
           

這時應該能看到__CUDA Information__裡找到了GPU裝置，然後像什麼cublas, cusparse, cufft, curand, nvvm這些庫應該也都找到了。這就配置好了，是不是很簡單？

最後隻需要在代碼中調用

os.environ['CUDA_VISIBLE_DEVICES']='0'

轉載來自這篇知乎