Как использовать autograph и tf.device с методом класса tf.function?

Question

В следующем коде для меня абсолютно необходимо выполнить полную функцию в графическом процессоре без единого перехода назад к центральному процессору.Это потому, что у меня 4 ядра процессора, но у меня 1200 ядер cuda.Теоретически это возможно, потому что тензорный поток feed_forwards, если операторы и назначения переменных могут быть выполнены на графическом процессоре (у меня NVIDIA GTX 1060).

Проблема, с которой я сталкиваюсь, состоит в том, что тензор потока 2.0 выполняет это автоматическое назначениедля GPU и CPU в бэкэнде и не упоминает, какие из его опций совместимы с GPU.Когда я запускаю следующую функцию с устройством в качестве графического процессора, я получаю

parallel_func could not be transformed and will be staged without change.

, и он последовательно запускается на графическом процессоре.

Мой вопрос: где использовать tf.device?Какая часть кода будет преобразована с помощью автографа в код графического процессора, а какая останется на процессоре?Как я могу преобразовать это тоже в 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

Answer 1

Примечание: сообщение parallel_func could not be transformed and will be staged without change выводится автографом, и поскольку оно содержит поток управления, зависящий от данных, вполне вероятно, что функция вообще не может быть запущена. Было бы целесообразно подать номер с инструкциями по воспроизведению и более подробными сообщениями журнала.