如何在带自定义Config类的Mask RCNN项目中配置TensorFlow GPU显存动态增长

你遇到的问题核心是两个config属于完全不同的配置体系，只是变量名撞了，修改变量名即可解决，具体操作如下：

解决步骤 #

1. 重命名TensorFlow会话配置的变量，避免和你项目里的训练配置变量config冲突，比如命名为tf_config
2. 显存动态增长的配置代码需要放在创建Mask RCNN模型之前执行，同时配套设置Keras全局会话保证配置生效

修改后的代码示例 #

首先在main.py头部补充导入：

import tensorflow as tf
# 如果你用的是TensorFlow自带的Keras，就换成 from tensorflow.keras import backend as K
from keras import backend as K

你原有逻辑部分修改为：

# Configurations
if args.command == "train":
    config = BalloonConfig()

# 新增TensorFlow显存动态增长配置，变量名用tf_config避免冲突
tf_config = tf.ConfigProto()
tf_config.gpu_options.allow_growth = True
session = tf.Session(config=tf_config)
K.set_session(session)

# Create model
if args.command == "train":
    model = modellib.MaskRCNN(mode="training", config=config,
                              model_dir=args.logs)

额外优化建议 #

如果配置后仍然存在显存不足问题，可以尝试调低BalloonConfig里的IMAGES_PER_GPU参数，或者缩小IMAGE_MAX_DIM、IMAGE_MIN_DIM的数值降低输入图片分辨率。

内容的提问来源于stack exchange，提问作者Jun Hong