Numba向量化在一个简单示例上无法正常工作

问题描述

我正在尝试为正在开发的程序添加gpu支持，并且遇到numba和cuda，然后遇到ufunc。我正在研究https://numba.pydata.org/numba-doc/dev/cuda/ufunc.html中的基本示例 我正在使用pycharm 2019.3.1 我正在使用numba 0.51.2

我得到了错误：

Traceback (most recent call last): line 19,in <module>
D = cu_discriminant(A,B,C)
TypeError: not enough arguments for format string

出于明显原因不共享路径 这是我的代码有问题还是numba问题？

这是我的代码：

import math
from numba import vectorize
import numpy as np
from time import perf_counter as count
Start = count()

@vectorize(['float32(float32,float32,float32)','float64(float64,float64,float64)'],target='cpu')
def cu_discriminant(a,b,c):
    return math.sqrt(b ** 2 - 4 * a * c)

N = 10000
dtype = np.float32

# prepare the input
A = np.array(np.random.sample(N),dtype=dtype)
B = np.array(np.random.sample(N) + 10,dtype=dtype)
C = np.array(np.random.sample(N),dtype=dtype)

D = cu_discriminant(A,C)

print(D)  # print result
Stop = count()
print(f'it took {round(Stop - Start,3)} seconds')

当target ='cpu'时有效，而当target ='cuda'时无效

我无所适从，因为GPU加速将大大加快我的代码的速度。任何帮助表示赞赏。

解决方法

您的代码在两种情况下都可以在没有“tensorflow 2”的情况下运行良好。

结果如下

在“target='cpu'”的情况下 [10.905523 10.944542 10.116297 ... 10.752323 10.669133 10.494761] 花了 0.083 秒

在“target='cuda'”的情况下 [10.605754 10.977879 10.845842 ... 10.56295 10.10678 10.774803] 花了 0.129 秒

# with rtx 3070

但是我发现安装了“tensorflow 2”的地方也出现了同样的错误。 错误信息是“NvvmError：编译失败”