如何将 Numba 用于 Pytorch 张量？

Question

我是 Numba 的新手，我需要使用 Numba 来加速一些 Pytorch 函数。但我发现即使是一个非常简单的功能也不起作用:(

import torch
import numba

@numba.njit()
def vec_add_odd_pos(a, b):
    res = 0.
    for pos in range(len(a)):
        if pos % 2 == 0:
            res += a[pos] + b[pos]
    return res

x = torch.tensor([3, 4, 5.])
y = torch.tensor([-2, 0, 1.])
z = vec_add_odd_pos(x, y)

但是出现如下错误 def vec_add_odd_pos(a, b): 资源 = 0。 ^

此错误可能是由以下参数引起的：

参数 0：无法确定
参数 1：无法确定

谁能帮帮我？也将不胜感激包含更多示例的 link。谢谢

Answer 1

numba 支持 numpy 数组但不支持 torch 的张量。然而有一座桥 Tensor.numpy():

Returns self tensor as a NumPy ndarray. This tensor and the returned ndarray share the same underlying storage. Changes to self tensor will be reflected in the ndarray and vice versa.

这意味着您必须将 jitted 函数调用为：

...
z = vec_add_odd_pos(x.numpy(), y.numpy())

如果 z 也应该是 torch.Tensor，那么 torch.from_numpy 就是我们需要的：

Creates a Tensor from a numpy.ndarray.

The returned tensor and ndarray share the same memory. Modifications to the tensor will be reflected in the ndarray and vice versa. The returned tensor is not resizable. ...

对于我们的代码，这意味着

...
z = torch.from_numpy(vec_add_odd_pos(x.numpy(), y.numpy()))

应该叫。

Answer 2

正如其他人提到的，numba 目前不支持 torch 张量，只支持 numpy 张量。但是有 TorchScript，它有类似的目标。然后可以将您的函数重写为：

import torch

@torch.jit.script
def vec_add_odd_pos(a, b):
    res = 0.
    for pos in range(len(a)):
        if pos % 2 == 0:
            res += a[pos] + b[pos]
    return res

x = torch.tensor([3, 4, 5.])
y = torch.tensor([-2, 0, 1.])
z = vec_add_odd_pos(x, y)

注意：虽然你说你的代码片段只是一个简单的例子，但 for 循环真的很慢运行 TorchScript 可能对你帮助不大，你应该不惜一切代价避免它们，只有在没有的时候才使用 then存在其他解决方案。话虽如此，以下是如何以更高效的方式实现您的功能：

def vec_add_odd_pos(a, b):
    evenids = torch.arange(len(a)) % 2 == 0
    return (a[evenids] + b[evenids]).sum()

Answer 3

Pytorch 现在在 GPU 张量上公开了一个接口，numba 可以直接使用它：

numba.cuda.as_cuda_array(tensor)

测试脚本提供了一些使用示例：https://github.com/pytorch/pytorch/blob/master/test/test_numba_integration.py

如何将 Numba 用于 Pytorch 张量？

How can I use Numba for Pytorch tensors?

numba

pytorch

tensor