Daniel Möller
Daniel Möller

Reputation: 86650

Tensorflow - Is it possible to manually decide which tensors in a graph go to the GPU and CPU?

I've created a network in Keras that uses both convolutional and LSTM layers.

And I've read that tensorflow can't handle well the LSTM layers in GPU: http://minimaxir.com/2017/07/cpu-or-gpu/

In fact, my network got slow in the GPU (hoping its not a further problem).

I want tensorflow to throw all convolutional operations to the GPU, but keep all the LSTM operations in the CPU. Is that possible?

So, considering I can find the tensorflow graph and identify each of its tensors, can I define manually where each tensor performs its operations?

Upvotes: 3

Views: 244

Answers (2)

aerin
aerin

Reputation: 22724

Look at this document: Distributed TensorFlow Specifying distributed devices in your model

Upvotes: 1

Marcin Możejko
Marcin Możejko

Reputation: 40526

Following implementation of this network, you could (if you use Tensorflow) try the following technique:

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

Upvotes: 2

Related Questions