RTX 4090训练Mask RCNN权重加载随机异常求助

问题根源 #

TensorFlow 2.x默认启用了更多异步、并行化运算，加上RTX 4090的Ampere架构特性，以及部分TF2.x版Mask RCNN移植代码的遗留随机逻辑，导致加载同一权重后出现输出不一致的情况；而TF1.x环境下的确定性控制逻辑更严格，GTX 1080 Ti不存在这类兼容性问题。

具体解决方案 #

1. 全局锁定所有随机种子 #

在代码最开头设置以下环境变量和种子，确保Python、NumPy、TensorFlow的随机行为完全可控：

import os
import numpy as np
import tensorflow as tf

# 锁定Python哈希种子
os.environ['PYTHONHASHSEED'] = '0'
# 启用TensorFlow确定性运算
os.environ['TF_DETERMINISTIC_OPS'] = '1'
os.environ['TF_CUDNN_DETERMINISTIC'] = '1'

# 设置NumPy种子
np.random.seed(42)
# 设置TensorFlow全局种子
tf.random.set_seed(42)
# 限制线程数避免并行运算引入随机性
tf.config.threading.set_inter_op_parallelism_threads(1)
tf.config.threading.set_intra_op_parallelism_threads(1)

2. 修正模型加载逻辑 #

加载权重时需避免额外的随机初始化操作，确保权重完全从文件读取：

• 加载模型时跳过编译阶段的随机初始化，之后手动控制训练模式参数：

# 以mrk1992的Mask RCNN实现为例
from mrcnn import model as modellib

model = modellib.MaskRCNN(mode="inference", config=config, model_dir=MODEL_DIR)
# 强制从权重文件读取所有参数，跳过不匹配检查
model.load_weights(WEIGHTS_PATH, by_name=True, skip_mismatch=False)

# 推理时必须显式设置training=False，避免BatchNorm/Dropout层的随机行为
results = model.detect([image], training=False)

• 检查Anchor Layer实现：确保Anchor生成逻辑在加载权重后不会重新随机生成，部分TF2.x移植版可能在模型初始化时动态生成Anchor，需修改代码让Anchor参数从权重文件读取，而非重新计算。

3. 验证权重一致性 #

加载权重后，打印某一关键层的权重哈希值，多次加载后对比是否一致，确认权重未被随机修改：

# 取RPN共享卷积层的权重示例
rpn_layer = model.keras_model.get_layer("rpn_conv_shared")
weights = rpn_layer.get_weights()[0]
# 打印权重的字节哈希值
print(hash(weights.tobytes()))

若多次加载后哈希值相同，说明权重加载正常，问题出在推理阶段的随机运算；若哈希值不同，需排查模型加载代码中是否存在覆盖权重的逻辑。

4. 修复TF2.x兼容性细节 #

• 替换所有numpy.bool为原生bool，避免版本兼容问题导致的隐式随机行为；
• 确保所有自定义层（如Anchor Layer、RoIAlign）的实现中，推理阶段完全禁用随机操作（如随机裁剪、随机Anchor生成等）。

验证步骤 #

1. 固定测试图像，多次加载同一权重后执行推理；
2. 对比每次推理的检测框坐标、掩码结果的哈希值；
3. 若结果一致，说明问题解决；若仍不一致，需排查是否有第三方依赖（如OpenCV图像读取）引入的随机性，可固定图像读取的参数（如cv2.imread的flags）。

内容的提问来源于stack exchange，提问作者xXAI-botXx