在TensorFlow中,GPU内存默认是一次性分配的,这意味着如果模型占用的内存超过可用内存的限制,将无法运行模型,而会出现OOM(Out Of Memory)错误。为了解决这个问题,TensorFlow提供了函数set_memory_growth,它可以让TensorFlow动态分配GPU内存,只使用所需的GPU内存。总之,使用set_memory_growth函数,可以在程序运行时分配所需的GPU内存,而不是在程序启动时将GPU内存分配给TensorFlow,这样可以避免在运行大型模型时出现内存不足的问题。
1 2 3
from tensorflow.config.experimental import list_physical_devices, set_memory_growth physical_devices = list_physical_devices('GPU') set_memory_growth(physical_devices[0], True)