预训练Keras模型VGG19训练时RMSprop优化器报KeyError的问题

问题描述 #

我使用以下代码加载ImageNet预训练VGG19模型并适配自定义数据集：

from keras.applications.vgg19 import VGG19


optim = tf.keras.optimizers.RMSprop(momentum=0.9)
vgg19 = VGG19(include_top=False, weights='imagenet', input_tensor=tf.keras.layers.Input(shape=(224, 224, 3)))
vgg19.trainable = False
# x = keras.layers.GlobalAveragePooling2D()(model_vgg19_pt.output)
x = keras.layers.Flatten()(vgg19.output)
output = keras.layers.Dense(n_classes, activation='softmax')(x)
model_vgg19_pt = keras.models.Model(inputs=[vgg19.input], outputs=[output])
model_vgg19_pt.compile(optimizer=optim,
                       loss='categorical_crossentropy', metrics=['categorical_accuracy'])
callback = tf.keras.callbacks.LearningRateScheduler(scheduler)
model_vgg19_pt.fit(x_train, y_train, batch_size=20,
                             epochs=50, callbacks=[callback]
                             )

执行model.fit()时触发以下错误：

KeyError: 'The optimizer cannot recognize variable dense_1/kernel:0. This usually means you are trying to call the optimizer to update different parts of the model separately. Please call optimizer.build(variables) with the full list of trainable variables before the training loop or use legacy optimizer `tf.keras.optimizers.legacy.{self.class.name}.'

补充信息 #

• 使用keras.applications.inception_v3时会出现同样问题
• 代码在TensorFlow CPU环境的Jupyter Notebook中可正常运行，但在TensorFlow-GPU的远程机器上报错
• 替换为SGD优化器无问题，仅RMSprop会报错
• 将优化器直接写在compile方法内即可正常运行，代码示例：

model_vgg19_pt.compile(optimizer=tf.keras.optimizers.RMSprop(momentum=0.9),
                           loss='categorical_crossentropy', metrics=['categorical_accuracy'])

疑问 #

1. 为什么只有RMSprop报错，SGD却没问题？
2. 为什么把优化器写在compile内部就可以正常运行？

问题分析与解答 #

1. RMSprop与SGD的差异原因 #

TensorFlow中不同优化器的变量绑定逻辑存在差异：

• RMSprop这类优化器需要维护额外的状态参数（比如动量项、滑动平均系数），必须在训练前明确绑定所有需要更新的可训练变量，才能初始化对应的状态张量。
• GPU环境下TensorFlow对优化器的变量校验更严格，提前实例化的RMSprop在模型构建完成后，没有感知到新增的自定义全连接层变量（如dense_1/kernel:0），因此训练时触发变量识别错误。
• SGD优化器的实现逻辑更简洁，内部兼容了延迟绑定变量的逻辑，即使提前实例化，也能在训练时自动绑定新增变量，因此不会报错。
• CPU环境下TensorFlow的校验规则相对宽松，允许优化器延迟绑定变量，所以代码可以正常运行。

2. 优化器写在compile内部的作用 #

当在compile方法内直接实例化优化器时，compile流程会自动完成以下操作：

• 此时模型已经完全构建，所有可训练变量（包括预训练模型的冻结变量、自定义层的可训练变量）都已生成。
• compile会调用优化器的build方法，传入模型的完整可训练变量列表，让优化器提前初始化所有需要的状态参数，确保训练时能识别所有待更新变量，因此不会触发错误。

可选解决方案 #

除了将优化器写在compile内部，还有两种可行方案：

• 手动绑定变量：在模型构建完成后，手动调用优化器的build方法，传入模型的可训练变量：

  model_vgg19_pt = keras.models.Model(inputs=[vgg19.input], outputs=[output])
  optim.build(model_vgg19_pt.trainable_variables)  # 手动绑定所有可训练变量
  model_vgg19_pt.compile(optimizer=optim, loss='categorical_crossentropy', metrics=['categorical_accuracy'])
  

• 使用Legacy优化器：切换到兼容旧逻辑的Legacy版本优化器，绕过严格的变量校验：

  optim = tf.keras.optimizers.legacy.RMSprop(momentum=0.9)
  

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