使用 numba 循环消耗 GPU 而不是 CPU答案

【问题标题】：Used numba to consume GPU instead of CPU in loop使用 numba 循环消耗 GPU 而不是 CPU
【发布时间】：2021-04-14 04:02:51
【问题描述】：

我有一个包含多个身份的列表，每个身份由多个图像组成。当我从 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])

【问题讨论】：

您的样本列表有多大？
尺寸为138368
帮助了吗 - stackoverflow.com/questions/18390336/…
实际上，上次我试图实现这个答案，但我坚持编码。请你帮助我好吗？上面的代码你能给出正确的答案吗？

标签： python python-3.x list cross-product

【解决方案1】：

Numba 目前正式支持支持 CUDA 的 GPU（AMD 支持仍处于试验阶段）。在 GPU 上使用 Numba 的官方文档/说明是 here。

根据您想要做什么，上面链接的网页中列出了各种详细信息，但如果您想让代码以简单的方式在 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

【讨论】：