tensorflow 显存 训练_TensorFlow训练模型,指定GPU训练,设置显存,查看gpu占用

1 linux查看当前服务器GPU占用情况:

nvidia-smi

周期性输出GPU使用情况: (如设置每3s显示一次GPU使用情况)

watch -n 3 nvidia-smi

效果如下:

2  指定GPU训练,使用CUDA_VISIBLE_DEVICES来指定

如果要指定第2块GPU训练,可以在python代码中如下指定:

import os

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

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

os.environ['CUDA_VISIBLE_DEVICES']='2，3'

当然也可以在运行python程序的时候进行指定:

CUDA_VISIBLE_DEVICES=2 python train_textcnn.py

3 设置定量的GPU使用量 和 设置最小的GPU使用量.

设置定量的GPU显存使用量:

config = tf.ConfigProto()

config.gpu_options.per_process_gpu_memory_fraction = 0.9 # 占用GPU90%的显存

session = tf.Session(config=config)

设置最小的GPU显存使用量:

config = tf.ConfigProto()

config.gpu_options.allow_growth = True

session = tf.Session(config=config)