В следующем коде мне абсолютно необходимо выполнить всю функцию на графическом процессоре без единого перехода обратно на ЦП. Это потому, что у меня 4 ядра процессора, но у меня 1200 ядер cuda. Теоретически это возможно, потому что тензорный поток feed_forwards, операторы if и и присваивание переменных можно выполнять на GPU (у меня NVIDIA GTX 1060).
Проблема, с которой я столкнулся, заключается в том, что tensorflow2.0 автоматически назначает GPU и CPU в бэкэнде и не упоминает, какие из его операций совместимы с GPU. Когда я запускаю следующую функцию с устройством в качестве графического процессора, я получаю
parallel_func could not be transformed and will be staged without change.
и он работает последовательно на GPU.
У меня вопрос, где использовать tf.device? Какая часть кода будет преобразована по автографу в код GPU, а какая останется на CPU? Как я могу преобразовать это тоже в GPU?
@tf.function
def parallel_func(self):
for i in tf.range(114): #want this parallel on GPU
for count in range(320): #want this sequential on GPU
retrivedValue = self.data[i][count]
if self.var[i]==1:
self.value[i] = retrievedValue # assigns, if else
elif self.var[i]==-1: # some links to class data through
self.value[i] = -retrivedValue # self.data, self.a and self.b
state = tf.reshape(tf.Variable([self.a[i], self.b[i][count]]), [-1,2])
if self.workerSwitch == False:
action = tf.math.argmax(self.feed_forward(i, count, state))
else:
action = tf.math.argmax(self.worker_feed_forward(i, count, state))
if (action==1 or action==-1):
self.actionCount +=1