在用TensorFlow进行深度学习时,容易遇到显存爆满的情况。
如下:
numpy.core._exceptions.MemoryError: Unable to allocate 359. MiB for an array with shape (60000, 28, 28) and data type float64
解决方案就是设置TensorFlow2设置显存按需分配
gpus = tf.config.experimental.list_physical_devices(device_type='GPU')
for gpu in gpus:
tf.config.experimental.set_memory_growth(gpu, True)