TensorFlow深度学习设置显存占用率

I. 两种方法查看显存使用情况

• 方法1：CMD窗口中输入命令nvidia-smi

• 方法2： 使用软件AIDA64（推荐），实时监测电脑硬件使用情况

电脑闲置可用显存为7443 MB。

II. 设置TensorFlow显存占用率

• 方法1：按需动态分配显存（更灵活）

config = tf.ConfigProto()
config.gpu_options.allow_growth = True # 按需动态分配显存
session = tf.Session(config=config)

• 方法2：指定显存使用比例

config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.8  # 占用80%显存
session = tf.Session(config=config)

III. 代码插入位置

修改两个.py文件，将代码插入.py文件的import modules行中。

• 文件1：model_main.py，路径：D:\tf_train\models\research\object_detection
• 文件2：trainer.py，路径：D:\tf_train\models\research\object_detection\legacy

IV. 使用对比

• 未使用显存配置，直接训练TensorFlow，显存占用7888 MB；

• 使用动态显存配置，显存占用5420 MB。

V. 参考资料

[1] https://stackoverflow.com/questions/34199233/how-to-prevent-tensorflow-from-allocating-the-totality-of-a-gpu-memory
[2] https://blog.youkuaiyun.com/u012856866/article/details/104921811