使用numba在循环中消耗GPU而不是CPU

Question

我有一个包含多个标识的列表，每个标识由多个图像组成。当我从json列表中检索正面图像时，它工作得很好。在那之后，我通过对成对形式的每个图像进行叉积来混合这个正列表，然后保存在负数组中。当我做交叉产品时，我的系统完全挂起了，即使我有16 GB的RAM和GPU。

代码

for i in range(0, len(idendities) - 1):
for j in range(i + 1, len(idendities)):
# print(samples_list[i], " vs ",samples_list[j])
cross_product = itertools.product(samples_list[i], samples_list[j])
cross_product = list(cross_product)
# print(cross_product)
for cross_sample in cross_product:
# print(cross_sample[0], " vs ", cross_sample[1])
negative = []
negative.append(cross_sample[0])
negative.append(cross_sample[1])
negatives.append(negative)
negatives = pd.DataFrame(negatives, columns=["file_x", "file_y"])
negatives["decision"] = "No"
negatives = negatives.sample(positives.shape[0])

Answer 1

Numba目前正式支持支持CUDA的GPU(AMD支持仍处于实验阶段(。在GPU上使用Numba的官方文档/说明在这里。

根据你到底想做什么，上面链接的网页中列出了各种细节，但如果你想让你的代码以一种简单的方式在GPU上运行，你可以做这样的事情：

from numba import cuda
@cuda.jit(device=True):
def foo(a: int, b: int, c: int):
return a * b + c #this can be whatever you want that is Numba compatible