如何修改我的 Python 函数才能使用 Numba，以便我可以在 Google Colab GPU/TPU 上加速它？

Question

我有一个Python函数，速度很慢，我相信如果在GPU/TPU上处理它会更快。我正在使用谷歌Colab。如何修改它以便我可以使用 Numba 在 Google Colab GPU 上处理它？

def Tournament (x, w, cutoff, y):
  preds = np.matmul(x,w)
  preds2 = (preds<np.quantile(preds, 1-cutoff,axis=0))
  preds = (preds>np.quantile(preds, cutoff,axis=0))
  rets = y * preds
  rets[rets == 0] = np.nan
  rets2 = y * preds2
  rets2[rets == 0] = np.nan
  ans = np.nanmedian(rets, axis=0) - np.nanmedian(rets2,axis=0)
  ans = ans[:,None]
  rand2 = np.random.uniform(-0.00000001,0.00000001,size=(popsize,1))
  ans += rand2
  garbage,sort = np.unique(ans,return_index = True)
  sort = sort[:,None]
  return sort, ans

我尝试加载 Numba 库并在函数之前使用

@jit

，但它似乎不起作用。

Answer 1

你可以使用

1.

@vectorize(device="cuda")

或2.

@guvectorize(input_data_type, input_array_shape, output_array_shape, device="cuda")

或3.

@numba.cuda.jit

第一个很容易实现，但性能提升有限，有时甚至更慢。

第二个需要对功能进行细微修改，我还没有真正得到它来加速我自己的功能，但它应该加速任何兼容的代码。

第三种需要对功能代码和主代码进行重大修改，您需要手动设计要复制到GPU设备和从GPU设备复制的数据流。这将为您提供最佳的 GPU 性能。