使用 GPU 版本的 TensorFlow 训练模型时,遇到 GPU 显存分配的错误

在使用 GPU 版本的 TensorFlow 训练模型时,有时候会遇到 GPU 显存分配的错误,
在这里插入图片描述
这是因为在调用 GPU 运行程序时,GPU 的显存空间不足而引起的。为了避免这个错误,可以在模型训练之前对 GPU 的使用模式进行设置。

这是列出当前系统中的所有 GPU ,放在列表 gpus 中,
在这里插入图片描述
我的主机中只有一块 GPU ,索引是0,把它设置为 memory_growth 模式,即允许内存增长。
在这里插入图片描述
也就是说,在程序运行的过程中,根据需要,为 TensorFlow 进程分配 GPU 显存,需要多少就分配多少,而不会过多的占用。如果系统中有多个 GPU , 可以使用循环语句,如下:
在这里插入图片描述
从而把所有 GPU都设置为内存增长模式。

猜你喜欢

转载自blog.csdn.net/xuechanba/article/details/124975045
今日推荐