我的 CNN 模型在我的 GPU 上使用了太多内存。 如何在我的 CPU 内存上托管一些张量？

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.