小编Nic*_*Maq的帖子

我正在尝试在 python 中进行强化学习。我使用的是 Tensorflow 2.1，我的机器有多个 GPU（带有 CUDA 10.2 驱动程序 440.59）。我正在使用 tf.device() 在我的 GPU 上分配操作。我没有使用 tf.distribute.Strategy。我正在构建我的模型：

    with tf.device(USE_GPU):
        model = build_keras_Seq()

其中 build_keras_Seq() 使用功能 AP 创建模型：

    model = tf.keras.Model(inputs=inputs, outputs=outputs)

我所有的输入都是与我的模型在同一 GPU 上分配的张量。

    with tf.device(USE_GPU):
        self.images_gpu = tf.zeros(shape=(1,IMG_HEIGHT,IMG_WIDTH), dtype=tf.int16) # (165, 160, 1)
        self.actions_gpu = tf.zeros(shape=(1,), dtype=tf.int16)
        self.rewards_gpu = tf.zeros(shape=(1,), dtype=tf.int16)
        self.dones_gpu = tf.zeros(shape=(1,), dtype=tf.int16)

我的目标是由 @tf.function 计算的，它实现了预期的 SARSA 并在 GPU 上返回一个张量：

    target_train is on device:  /job:localhost/replica:0/task:0/device:GPU:1

当我调用model.fit时，似乎在CPU上执行了很多操作（见下文）导致性能不佳。我的理解是张量在被发送到 GPU 之前被移回 CPU。关于如何防止这种行为并将张量直接从 GPU 馈送到同一 GPU 上托管的模型的任何想法？

2020-02-23 09:49:32.100259: I tensorflow/core/common_runtime/eager/execute.cc:573] Executing op …