My CNN Model uses too much memory on my GPU. How can I host some Tensors on my CPU memory?

Question

I'm training a CNN model on an NVidia RTX 2080, it became bigger and bigger and I have now some memory issues with the card. I read some paper about this subject, and it seems possible with Tensorflow to host some nodes on the CPU memory during the training and retrieve it in the GPU memory later when needed (as in http://learningsys.org/nips17/assets/papers/paper_18.pdf ).

Any ideas/docs/examples?

Thanks!

Answer 1

Without any code it's difficult to help. Generally you can take a look at the documentation .

For example with:

with tf.device('/cpu:0'):
  a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[2, 3], name='a')
  b = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[3, 2], name='b')

you can create variables explicitly on the cpu.