python - 如何将 autograph 和 tf.device 与 tf.function 包装的类方法一起使用？

Question

在下面的代码中，我绝对必须在 GPU 中执行完整的函数，而不需要一次跳转回 CPU。这是因为我有 4 个 CPU 核心，但我有 1200 个 cuda 核心。理论上，这是可能的，因为 tensorflow feed_forwards、if 语句和变量分配可以在 GPU 上完成(我有 NVIDIA GTX 1060)。

我面临的问题是tensorflow2.0在后端自动分配GPU和CPU，并且没有提及其中哪些操作与GPU兼容。当我使用设备作为 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

Answer 1

旁注:消息parallel_func无法转换，将不做更改地暂存是通过签名输出的，并且由于它包含依赖于数据的控制流，因此该函数可能无法运行根本不。值得提交issue包含重现步骤和更详细的日志消息。