使用Keras Tuner+MirroredStrategy恢复检查点报错问题

问题背景 #

使用Keras Tuner为处理EMG数据的CNN进行超参数调优，基于TensorFlow的MirroredStrategy在4个NVIDIA GPU和5个CPU的环境运行。调参器可完成所有trial的2个epoch训练，但Hyperband下一次迭代尝试恢复已保存模型以继续训练时，抛出错误：

Make sure the slot variables are created under the same strategy scope. This may happen if you're restoring from a checkpoint outside the scope

错误原因 #

核心问题是模型构建（包括优化器的slot变量）与checkpoint恢复时的策略作用域不匹配。虽然代码中给Tuner指定了distribution_strategy=strategy，但HyperModel的build方法内的模型构建、编译逻辑并未包裹在strategy.scope()上下文管理器中，导致训练时生成的checkpoint里，优化器的slot变量（如Adam优化器的动量、方差变量）并非在分布式策略作用域下创建，后续恢复时触发作用域不匹配的错误。

解决方案 #

修改CNNHyperModel的build方法，将模型构建、编译的全部逻辑放到strategy.scope()内，确保所有变量（包括优化器的slot变量）都在分布式策略作用域下创建。

修改后的完整代码 #

class CNNHyperModel(HyperModel):

  def __init__(self, input_shape, output_shape):
    self.input_shape = input_shape
    self.output_shape = output_shape

  def build(self, hp):
    # 将模型构建与编译逻辑包裹在策略作用域内
    with strategy.scope():
      visible = Input(shape = (self.input_shape[0], self.input_shape[1], 1))
      
      filter_number = hp.Int("filter_base_size", 5, 7)

      conv = Conv2D(filters=2**(filter_number), \
        kernel_size=(4,4), activation='relu', padding = 'same')(visible)
      pooling = MaxPooling2D(pool_size=(2, 2), strides=2, padding='valid')(conv)

      conv2 = Conv2D(filters=2**(filter_number + 1), \
        kernel_size=(3,3), activation='relu', padding = 'same')(pooling)
      pooling2 =  MaxPooling2D(pool_size=(2, 2), strides=2, padding='valid')(conv2)

      conv3 = Conv2D(filters=2**(filter_number + 2), \
        kernel_size=(2,2), activation='relu', padding = 'same')(pooling2)
      pooling3 =  MaxPooling2D(pool_size=(2, 2), strides=2, padding='valid')(conv3)

      flatten = Flatten()(pooling3)
      dropout = Dropout(0.5)(flatten)
      
      num_layers = hp.Int("num_layers", 3, 4)
      dense_units = [hp.Int("dense_units_1", 80, 240, 30), \
          hp.Int("dense_units_2", 80, 240, 30), hp.Int("dense_units_3", 80, 240, 30), \
          hp.Int("dense_units_4", 80, 240, 30)]
      hidden = dropout

      for i in range(num_layers):
        hidden = Dense(units=dense_units[i], activation = 'relu')(hidden)

      output = Dense(units=self.output_shape[0], activation = 'softmax')(hidden)

      model = Model(inputs=visible, outputs=output)

      model.compile(
          optimizer="adam",
          loss="categorical_crossentropy",
          metrics = "accuracy",
      )
    return model
    
    
strategy = tf.distribute.MirroredStrategy()
batch_size = 800
epochs = 80
val_dataset = (np.array(testing_input), np.array(testing_output))


tuner = Hyperband(
  hypermodel = CNNHyperModel(
    input_shape = [len(training_input[0]), len(training_input[0][0])],
    output_shape = [len(training_output[0])],
  ),
  objective='accuracy',
  distribution_strategy=strategy,
  hyperband_iterations=5
)

early_stopping_callback = keras.callbacks.EarlyStopping(
  monitor='val_loss', 
  patience=5
)

tuner.search(
  x = np.array(training_input),
  y = np.array(training_output),
  validation_data=val_dataset,
  callbacks=[early_stopping_callback],
  batch_size=batch_size,
  epochs = epochs,
)

best_model = tuner.get_best_models(num_models=1)[0]

# 加载最优模型后训练也需要包裹在策略作用域内
with strategy.scope():
  best_model.fit(
      x = np.array(training_input),
      y = np.array(training_output),
      epochs=epochs,
      batch_size=batch_size,
      validation_data=val_dataset,
  )

额外注意事项 #

• 确保strategy变量在CNNHyperModel类定义前初始化，保证build方法能正确访问到策略实例
• 所有涉及模型创建、编译、训练的逻辑，只要涉及分布式策略，都需要放到strategy.scope()内，避免变量作用域不匹配问题

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