GPU 内存处理
当您开始运行 TensorFlow 会话时,默认情况下它会抓取所有 GPU 内存,即使您将操作和变量仅放置在多 GPU 系统中的一个 GPU 上也是如此。如果您尝试同时运行另一个会话,则会出现内存不足错误。这可以通过多种方式解决:
- 对于多 GPU 系统,请设置环境变量
CUDA_VISIBLE_DEVICES=<list of device idx>
os.environ['CUDA_VISIBLE_DEVICES']='0'
在此设置之后执行的代码将能够获取仅可见 GPU 的所有内存。
- 当您不希望会话占用 GPU 的所有内存时,您可以使用配置选项
per_process_gpu_memory_fraction
来分配一定百分比的内存:
config.gpu_options.per_process_gpu_memory_fraction = 0.5
这将分配所有 GPU 设备的 50%的内存。
- 您还可以结合上述两种策略,即只制作一个百分比,同时只让部分 GPU 对流程可见。
- 您还可以将 TensorFlow 进程限制为仅在进程开始时获取所需的最小内存。随着进程的进一步执行,您可以设置配置选项以允许此内存的增长。
config.gpu_options.allow_growth = True
此选项仅允许分配的内存增长,但内存永远不会释放。
您将学习在后面的章节中跨多个计算设备和多个节点分配计算的技术。