Tensorflow占满显存怎么办？

Tensorflow默认分配设备的所有内存，如果想配置按需使用，可以如下操作，

1
2
3
4
5
import tensorflow as tf

physical_devices = tf.config.experimental.list_physical_devices('GPU')
assert len(physical_devices) > 0, "Not enough GPU hardware devices available"
tf.config.experimental.set_memory_growth(physical_devices[0], True)

以上把第一颗GPU设置为按需分配，其他的GPU也类似的设置。

转载请包括本文地址：https://allenwind.github.io/blog/12320
更多文章请参考：https://allenwind.github.io/blog/archives/