colab 使用 numba

原创已于 2022-05-08 15:04:33 修改 · 326 阅读

0 ·

CC 4.0 BY-SA版权

文章标签：

#深度学习 #python #pytorch

于 2022-05-08 11:45:51 首次发布

本文介绍了如何在Google Colab环境中安装NVIDIA CUDA工具包，配置环境变量，并通过Numba实现CUDA并行计算。通过实例演示了使用@cuda.jit装饰器编写和运行简单的Hello World程序，以展示Numba在GPU计算中的优势。

colab 使用 numba

必须安装cudatoolkit

关键问题在colab要勾选gpu选项

!apt-get install nvidia-cuda-toolkit

!apt-get install nvidia-cuda-toolkit
!pip3 install numba

import os
os.environ['NUMBAPRO_LIBDEVICE'] = "/usr/lib/nvidia-cuda-toolkit/libdevice"
os.environ['NUMBAPRO_NVVM'] = "/usr/lib/x86_64-linux-gnu/libnvvm.so"

from numba import cuda
import numpy as np
import time

@cuda.jit
def hello(data):
    data[cuda.blockIdx.x, cuda.threadIdx.x] = cuda.blockIdx.x

numBlocks = 5
threadsPerBlock = 10

data = np.ones((numBlocks, threadsPerBlock), dtype=np.uint8)

hello[numBlocks, threadsPerBlock](data)

print(data)