如何使用 numba/CUDA 加速 python 对象函数？

Question

I'm new to CUDA(installed numba about an hour ago). I'd like to speed up this function that's inside a class.

def predict(self, X):
   num_test = X.shape[0]
   Ypred = np.zeros(num_test, dtype=self.ytr.dtype)

   for i in range(num_test):
      distances = np.sum(np.abs(self.Xtr-X[i, :]), axis=1)
      min_index = np.argmin(distances)
      Ypred[i] = self.ytr[min_index]
      print(i)

   return Ypred

X is a 2D array of type float32 and Ypred is an array of type int32. I tried to speed it by inserting the following line just above the function.

  @vectorize(['int32(float32)'], target='cuda')

This gave me a huge list of errors but the significant part of it appears to be:

TypeError: Failed at nopython (analyzing bytecode)
Signature mismatch: 1 argument types given, but function takes 2 arguments

While I know exactly what the error says, I've no idea how to fix it. So... how do I make it work? Thanks in advance.

UPDATE:

I should've done a proper google search before asking(I did search, but I used the word 'object' instead of 'class', which gave me no useful results). The documentation helped me a lot, but now I've got these errors in my face and I've no clue what to do.

numba.errors.LoweringError: Failed at nopython (nopython mode backend)
Can only insert float* at [4] in {i8*, i8*, i64, i64, float*, [2 x i64], 
[2 x i64]}: got double*
File "main.py", line 40
[1] During: lowering "(self).Xtr = X" at D:/myStuff/DL/Week 3/1/main.py (40)
[2] During: resolving callee type: 
BoundFunction((<class 'numba.types.misc.ClassInstanceType'>, 'train')       f   for instance.jitclass.NearestNeighbours#24f01184f58<Xtr:array(float32, 2d, A),ytr:array(int32, 1d, A)>)
[3] During: typing of call at <string> (3)
--%<-----------------------------------------------------------------

File "<string>", line 3

Here's the entire class in its current state:

spec = [("Xtr", float32[:, :]), ("ytr", int32[:])]

@jitclass(spec)
class NearestNeighbours(object):
  def __init__(self):
    pass

  def train(self, X, y):
    self.Xtr = X            #line 40
    self.ytr = y

  def predict(self, X):
    num_test = X.shape[0]
    Ypred = np.zeros(num_test, dtype=self.ytr.dtype)

    for i in range(num_test):
      distances = np.sum(np.abs(self.Xtr-X[i, :]), axis=1)
      min_index = np.argmin(distances)
      Ypred[i] = self.ytr[min_index]
      print(i)

    return Ypred

UPDATE 2: Gave up on jitting the class and tried linking predict to its external clone. Using an empty jit seemed to work, but linking to cuda(for the speed) led to all sorts of fancy errors. I'll give it a rest for today and answer my own question if I somehow manage to resolve. Until a few hours ago, I thought that GPU acceleration would be as simple as adding an additional library or switching to a different compiler or something... but man... I had no idea I'd be in for such a bumpy ride.

Answer 1

As far as I can see your function only relies on X so there is no reason to have it as function in a class. Either declare it static @staticmethod (link) or take it out of the class - scope.

Presto: only 1 function parameter left.