How can I use Numba for Pytorch tensors?

user_2669433 0
Viewed 7515

I am new to Numba and I need to use Numba to speed up some Pytorch functions. But I find even a very simple function does not work :(

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)

But the following error appears def vec_add_odd_pos(a, b): res = 0. ^

This error may have been caused by the following argument(s):

  • argument 0: cannot determine Numba type of <class 'torch.Tensor'>
  • argument 1: cannot determine Numba type of <class 'torch.Tensor'>

Can anyone help me? A link with more examples would be also appreciated. Thanks.

numba tensor pytorch
3 Answers
user_57864750

As others have mentioned, numba currently doesn't support torch tensors, only numpy tensors. However there is TorchScript, which has a similar goal. Your function can then be rewritten as such:

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)

Beware: although you said your code snippet was just a simple example, for loops are really slow and running TorchScript might not help you much, you should avoid them at any cost and only use then when no other solution exist. That being said, here's how to implement your function in a more performant way:

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

numba supports numpy-arrays but not torch's tensors. There is however a bridge 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.

That means you have to call jitted functions as:

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

If z should be a torch.Tensor as well, torch.from_numpy is what we need:

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. ...

For our code that means

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

should be called.

user_57694630
Related