Я пишу очень специфическую нейронную сеть, и у меня есть много классов различных функций активации, каждая из которых имеет функцию для обычного python и одну функцию устройства. Проблема заключается в вызове этого метода из ядра CUDA.
@cuda.jit(device=True)
def activation_fn(z):
return max(0, z)
@cuda.jit
def backprop_kernel(arr):
arr[cuda.threadIdx.x] = activation_fn(arr[cuda.threadIdx.x])
def backprop_GPU(x, y):
arr = np.array([-3, -2, -1, 0, 1, 2, 3])
print(arr)
backprop_kernel[1, 7](arr)
print(arr)
backprop_GPU(None, None)
Это прекрасно работает, но я хочу, чтобы приведенный ниже код работал.
class Activation:
@cuda.jit(device=True)
def fn(z):
return max(0, z)
class Network:
def __init__(self):
self.activation_fn = Activation()
@cuda.jit
def kernel(arr):
arr[cuda.threadIdx.x] = activation_fn(arr[cuda.threadIdx.x])
def backprop(self, x, y):
arr = np.array([-3, -2, -1, 0, 1, 2, 3])
self.kernel[1, 7](arr)
net = Network()
net.backprop(None, None)
Как сделать «activation_fn» доступным из ядра?
@cuda.jit
должен использоваться с функциями, а не членами, поэтому вам нужно определить оформленные функции внутри методов и захватить функцию активации при определении ядра:
from numba import cuda
import numpy as np
class Activation:
def __init__(self):
@cuda.jit(device=True)
def fn(z):
return max(0, z)
self.fn = fn
class Network:
def __init__(self):
self.activation = Activation()
activation_fn = self.activation.fn
@cuda.jit
def kernel(arr):
arr[cuda.threadIdx.x] = activation_fn(arr[cuda.threadIdx.x])
self.kernel = kernel
def backprop(self, x, y):
arr = np.array([-3, -2, -1, 0, 1, 2, 3])
self.kernel[1, 7](arr)
print(arr)
net = Network()
net.backprop(None, None)
печатает:
$ python repro.py
[0 0 0 0 1 2 3]
(Обратите внимание, я пропустил предупреждения о производительности, которые появляются здесь, поскольку они ортогональны рассматриваемой проблеме)