欢迎您访问程序员文章站本站旨在为大家提供分享程序员计算机编程知识!
您现在的位置是: 首页

TensorFlow 指定 GPU 训练模型

程序员文章站 2022-07-13 13:18:19
...

Linux 查看当前服务器 GPU 的占用情况可以使用 nvidia-smi 命令,如下所示:

nvidia-smi

TensorFlow 指定 GPU 训练模型

关于 nvidia-smi 命令输出的详细解释,可参考笔者的另外一篇博客:GPU状态监测 nvidia-smi 命令详解。在此不再赘述,本文主要分享一下在用 TensorFlow 训练模型时如何指定 GPU。

在用 TensorFlow 训练深度学习模型的时候,若无手动指定,默认是选用第0块来训练,而且其他几块 GPU 也会被显示占用。有时候,我们更希望可以自己指定一块或者多块 GPU 来训练模型,接下来介绍一种常用的也是比较有效的方法:

举个例子,如果要指定只用第0块 GPU 训练,可以在 python 代码中如下指定:

import os
os.environ['CUDA_VISIBLE_DEVICES']='0'

运行 python 程序后,可发现就只有第0块 GPU 显示被占用,如上图所示。

如果要指定多块 GPU,可以如下:

os.environ['CUDA_VISIBLE_DEVICES']='1,2'

当然,上述这种方法不太方便,每次运行若换用 GPU 训练都需要改代码,可以在运行 python 程序的时候进行指定:

CUDA_VISIBLE_DEVICES=0 python textCnn.py