使用TensorFlow Lite Model Maker训练目标检测模型如何中断后恢复训练

前置配置（训练启动前设置，避免后续无法恢复） #

TensorFlow Lite Model Maker的目标检测训练接口原生支持检查点自动保存，只需要在调用训练方法时传入对应的配置参数即可，修改后的训练代码示例如下：

from tflite_model_maker import object_detector

# 原有数据集加载逻辑保持不变
train_data, val_data, test_data = # 你的数据集加载代码

# 初始化模型时添加检查点配置
model = object_detector.create(
    train_data,
    model_spec=object_detector.EfficientDetLite0Spec(), # 替换为你实际使用的模型规格
    validation_data=val_data,
    epochs=50,
    batch_size=32,
    checkpoint_dir='./tf_lite_obj_det_ckpts', # 检查点本地保存路径，*确保路径固定不要随意修改*
    checkpoint_freq=1, # 每1个epoch保存一次检查点，可根据需求调整频率
    do_train=False
)

# 启动训练
model.train(
    epochs=50,
    batch_size=32,
    checkpoint_dir='./tf_lite_obj_det_ckpts',
    checkpoint_freq=1
)

配置完成后，程序会在每个epoch训练结束后自动保存当前的模型权重、优化器状态、训练进度等信息到指定的文件夹中，无需手动操作。

中途停止训练操作 #

• 正常主动停止：直接在Jupyter Notebook工具栏点击「中断内核」即可停止训练，最新的训练进度已经自动保存
• 意外断电/程序崩溃：不会丢失最后一次成功保存的检查点数据

恢复训练操作 #

下次需要继续训练时，按照如下步骤操作即可从上次中断的epoch继续训练：

from tflite_model_maker import object_detector
import tensorflow as tf

# 1. 加载和第一次训练完全一致的数据集，不要修改数据集划分、预处理逻辑
train_data, val_data, test_data = # 原有数据集加载代码

# 2. 初始化和第一次训练完全相同配置的模型
model_spec = object_detector.EfficientDetLite0Spec() # 和第一次使用的模型规格保持一致
model = object_detector.create(
    train_data,
    model_spec=model_spec,
    validation_data=val_data,
    epochs=50,
    batch_size=32,
    do_train=False # 禁止自动启动训练
)

# 3. 加载最新的检查点文件
latest_checkpoint = tf.train.latest_checkpoint('./tf_lite_obj_det_ckpts')
model.model.load_weights(latest_checkpoint)

# 4. 计算剩余需要训练的epoch数
# 可通过检查点文件名判断已完成的epoch数，例如ckpt-10对应已完成10个epoch
finished_epoch_num = 10
remaining_epoch_num = 50 - finished_epoch_num

# 5. 启动续训，保留检查点配置
model.train(
    epochs=remaining_epoch_num,
    batch_size=32,
    checkpoint_dir='./tf_lite_obj_det_ckpts',
    checkpoint_freq=1
)

注意：两次训练的模型规格、batch size、数据集配置必须完全一致，否则会出现张量维度不匹配的问题，导致检查点加载失败。如果需要修改以上参数，需要从零开始训练。

内容的提问来源于stack exchange，提问作者Uğur Erçelik