将训练好的权重保存在机器学习代码中

Question

我有用于运行机器学习模型的colab，但是当它达到80个epoch时，我的colab内存崩溃了，我无法超越80个epoch。我希望有人帮助我将训练后的权重保存在某个地方，在内存崩溃后，我开始训练来自那个时代的模型。这是我的代码，我如何在这个 python 代码中编写该目的代码以及在哪里？

for comm_round in range(comms_round):

    global_weights = global_model.get_weights()

    scaled_local_weight_list = list()

    client_names= list(clients_batched.keys())
    random.shuffle(client_names)

    for client in client_names:
        local_model = Transformer
        local_model.compile(loss=tf.keras.losses.CategoricalCrossentropy(),
                            optimizer=tf.keras.optimizers.Adam(learning_rate = 0.001),
                            metrics='acc')

        local_model.set_weights(global_weights)

        local_model.fit(clients_batched[client], epochs=1, verbose=0, callbacks=[checkpoint_callback])

        scaling_factor = weight_scalling_factor(clients_batched, client)
        scaled_weights = scale_model_weights(local_model.get_weights(), scaling_factor)
        scaled_local_weight_list.append(scaled_weights)

        K.clear_session()

    average_weights = sum_scaled_weights(scaled_local_weight_list)

    global_model.set_weights(average_weights)

    for(X_test, Y_test) in test_batched:
        global_acc, global_loss = test_model(test_x, test_y, global_model, comm_round + 1)

此代码用于最后一步和联邦学习。

Answer 1

潜在的解决方案

---

注意：我发现你的

checkpoint_callback

函数中有

fit

，我认为这可能是模型回调？如果是这样，不确定这是否有效，但这是我在使用 TensorFlow 框架的训练期间尝试保存最佳权重时想到的通用方法。由于这个过程在各个时期都是连续的，即使它在晚期时期失败，我想应该已经保存了之前的最佳运行。

您可以使用 TensorFlow 的

ModelCheckpoint

回调在每个 epoch 后自动保存模型的权重。如果运行时崩溃，您可以加载这些权重以恢复训练。设置方法如下：

1. 在训练循环之前定义

   ModelCheckpoint

   回调：

checkpoint_filepath = '/path/to/checkpoint/directory'
checkpoint_callback = tf.keras.callbacks.ModelCheckpoint(
    filepath=checkpoint_filepath,
    save_weights_only=True,
    monitor='val_loss',  # change this to what you want to monitor
    mode='auto', 
    save_best_only=False,
    verbose=1,
)

2. 在您的

   ModelCheckpoint

   函数中包含

   fit

   回调：

for comm_round in range(comms_round):
    # ... existing code ...

    for client in client_names:
        # ... existing code ...

        local_model.fit(clients_batched[client], epochs=1, verbose=0, callbacks=[checkpoint_callback])

        # ... existing code ...

每个 epoch 之后，模型的权重都会保存到指定的文件路径。

3. 如果您的运行时崩溃并且需要恢复训练，您可以将保存的权重加载到模型中：

model.load_weights(checkpoint_filepath)

请注意，保存模型权重的目录路径必须存在。如果您使用 Google Colab，将您的权重保存在 Google 云端硬盘中可能是个好主意。这样，即使 Colab 运行时被回收，保存的权重也将持续存在。为此，您需要将 Google Drive 安装到 Colab 笔记本。

另请注意，这不会自动保存当前的通信回合。您需要手动管理它。例如，您可以在每次保存模型权重时将当前通信轮数保存到文件中。当您恢复训练时，您可以从文件中读取此数字并从那里继续进行通信。