Вызовите Model.fit с тензорами / Выполнение операции на ЦП, а не ЦП / Tensorflow 2.1 - PullRequest
Новая
       40

Вызовите Model.fit с тензорами / Выполнение операции на ЦП, а не ЦП / Tensorflow 2.1

1 голос
/ 23 февраля 2020

Я экспериментирую с обучением подкреплению в python. Я использую Tensorflow 2.1, и на моей машине установлены графические процессоры muyliple (с драйвером CUDA 10.2 440.59). Я распределяю операции на моих графических процессорах, используя tf.device (). Я не использую tf.distribute.Strategy. Я строю свою модель с: 

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

, где build_keras_Seq () использует функциональную точку доступа для создания модели: 

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

Все мои входные данные являются тензорами, расположенными на том же графическом процессоре, что и мой модель. 

    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.f, которая реализует ожидаемую SARSA и возвращает тензор на графическом процессоре: 

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

Когда я вызываю модель.fit, она Кажется, что много операций выполняется на процессоре (см. ниже), что приводит к снижению производительности. Насколько я понимаю, тензоры передаются обратно в ЦП перед отправкой в ​​ГП. Любая идея о том, как предотвратить такое поведение и передать тензоры непосредственно из графического процессора в модель, размещенную на том же графическом процессоре? 

2020-02-23 09:49:32.100259: I tensorflow/core/common_runtime/eager/execute.cc:573] Executing op RangeDataset in device /job:localhost/replica:0/task:0/device:CPU:0
2020-02-23 09:49:32.101114: I tensorflow/core/common_runtime/eager/execute.cc:573] Executing op RepeatDataset in device /job:localhost/replica:0/task:0/device:CPU:0
2020-02-23 09:49:32.108407: I tensorflow/core/common_runtime/eager/execute.cc:573] Executing op MapDataset in device /job:localhost/replica:0/task:0/device:CPU:0
2020-02-23 09:49:32.109087: I tensorflow/core/common_runtime/eager/execute.cc:573] Executing op PrefetchDataset in device /job:localhost/replica:0/task:0/device:GPU:1
2020-02-23 09:49:32.117795: I tensorflow/core/common_runtime/eager/execute.cc:573] Executing op FlatMapDataset in device /job:localhost/replica:0/task:0/device:CPU:0
2020-02-23 09:49:32.118524: I tensorflow/core/common_runtime/eager/execute.cc:573] Executing op TensorDataset in device /job:localhost/replica:0/task:0/device:CPU:0
2020-02-23 09:49:32.119764: I tensorflow/core/common_runtime/eager/execute.cc:573] Executing op RepeatDataset in device /job:localhost/replica:0/task:0/device:CPU:0
2020-02-23 09:49:32.120133: I tensorflow/core/common_runtime/eager/execute.cc:573] Executing op ZipDataset in device /job:localhost/replica:0/task:0/device:CPU:0
2020-02-23 09:49:32.128411: I tensorflow/core/common_runtime/eager/execute.cc:573] Executing op ParallelMapDataset in device /job:localhost/replica:0/task:0/device:CPU:0
2020-02-23 09:49:32.129839: I tensorflow/core/common_runtime/eager/execute.cc:573] Executing op ModelDataset in device /job:localhost/replica:0/task:0/device:CPU:0

Спасибо!

N

...