Tensorflow - Est-il possible de décider manuellement quels tenseurs dans un graphique vont au GPU et au CPU?

Question

J'ai créé un réseau dans Keras qui utilise à la fois des couches convolutives et LSTM.

Et j'ai lu que tensorflow ne peut pas bien gérer les couches de LSMC en GPU: http://minimaxir.com/2017/07/cpu-or-gpu/

En fait, mon réseau obtenu lent dans le GPU (en espérant ne est pas un autre problème).

Je veux que tensorflow rejette toutes les opérations de convolution vers le GPU, mais conserve toutes les opérations LSTM dans la CPU. Est-ce possible? Donc, étant donné que je peux trouver le graphique tensorflow et identifier chacun de ses tenseurs, puis-je définir manuellement où chaque tenseur effectue ses opérations?

Answer 1

Après la mise en œuvre du réseau this, vous pouvez (si vous utilisez Tensorflow) essayez la technique suivante:

with tf.device('/cpu:0'): 
    # model layers which should be run on CPU go here 

Answer 2

Regardez ce document: Distributed TensorFlow Specifying distributed devices in your model