TensorFlow Model Maker训练EfficientDet后如何绘制损失曲线

一、搞定损失曲线绘制 #

因为tflite_model_maker的object_detector.create()不直接返回训练history，手动用TensorFlow回调抓损失就行：

• 自定义回调记录每轮损失

import tensorflow as tf
import matplotlib.pyplot as plt

# 写个回调类存损失
class LossLogger(tf.keras.callbacks.Callback):
    def __init__(self):
        super().__init__()
        self.train_losses = []
        self.val_losses = []

    def on_epoch_end(self, epoch, logs=None):
        self.train_losses.append(logs.get('loss'))
        self.val_losses.append(logs.get('val_loss'))
        print(f"第{epoch+1}轮：训练损失={logs.get('loss'):.4f}，验证损失={logs.get('val_loss'):.4f}")

loss_logger = LossLogger()

• 训练时带上这个回调
  修改你的训练代码，加个callbacks参数：

model = object_detector.create(
    train_data, 
    model_spec=spec, 
    batch_size=5, 
    train_whole_model=True, 
    epochs=28, 
    validation_data=val_data,
    callbacks=[loss_logger]  # 把回调传进去
)

• 训练完直接画曲线

epochs = range(1, 29)
plt.plot(epochs, loss_logger.train_losses, 'b-', label='训练损失')
plt.plot(epochs, loss_logger.val_losses, 'r-', label='验证损失')
plt.title('训练/验证损失变化')
plt.xlabel('轮次')
plt.ylabel('损失值')
plt.legend()
plt.show()

嫌自定义麻烦？直接用CSVLogger把日志存成CSV，之后读文件绘图也行：

csv_logger = tf.keras.callbacks.CSVLogger('training_logs.csv')
# 训练时传callbacks=[csv_logger]

# 事后读文件绘图
import pandas as pd
logs_df = pd.read_csv('training_logs.csv')
logs_df.plot(x='epoch', y=['loss', 'val_loss'], kind='line')
plt.show()

二、解决Colab断开重训+过拟合问题 #

1. 自动保存最优模型 #

用ModelCheckpoint回调，每轮自动存模型，还能只存验证损失最低的版本：

checkpoint_callback = tf.keras.callbacks.ModelCheckpoint(
    filepath='best_model.h5',  # 存到当前目录，要存Drive就改路径
    monitor='val_loss',
    save_best_only=True,  # 只留最好的模型
    verbose=1
)

# 训练时同时带损失日志和保存回调
model = object_detector.create(
    train_data, 
    model_spec=spec, 
    batch_size=5, 
    train_whole_model=True, 
    epochs=28, 
    validation_data=val_data,
    callbacks=[loss_logger, checkpoint_callback]
)

要存到Google Drive的话，先挂载：

from google.colab import drive
drive.mount('/content/drive')
# 把filepath改成'/content/drive/MyDrive/best_model.h5'

2. 从保存的模型继续训练 #

断开后不用从头训，加载已存的模型接着来：

from tflite_model_maker import object_detector

# 先重新加载数据和模型规格
train_data = object_detector.DataLoader.from_pascal_voc('fire4/training','fire4/training', label_map=['fire_exit'])
val_data = object_detector.DataLoader.from_pascal_voc('fire4/validation', 'fire4/validation', label_map=['fire_exit'])
spec = object_detector.EfficientDetSpec(...)  # 重新定义你用的模型规格

# 加载保存的模型
model = object_detector.ObjectDetector(
    model_spec=spec,
    train_data=train_data,
    validation_data=val_data
)
model.load_weights('best_model.h5')  # 换成你的保存路径

# 继续训练，epochs设为剩余轮次（比如已经训了10轮，就设18）
model.train(
    batch_size=5,
    train_whole_model=True,
    epochs=18,
    callbacks=[loss_logger, checkpoint_callback]
)

3. 避免过拟合的小技巧 #

• 加EarlyStopping回调，验证损失不下降就自动停：

early_stopping = tf.keras.callbacks.EarlyStopping(
    monitor='val_loss',
    patience=5,  # 5轮没提升就停
    restore_best_weights=True  # 自动回到最优权重
)
# 加到callbacks列表里

• 分阶段训练：先冻结backbone只训头部，再微调全模型，减少过拟合：

# 第一步：只训头部
model = object_detector.create(
    train_data,
    model_spec=spec,
    batch_size=5,
    train_whole_model=False,
    epochs=10,
    validation_data=val_data,
    callbacks=[loss_logger, checkpoint_callback]
)
# 第二步：微调全模型
model.train(
    batch_size=5,
    train_whole_model=True,
    epochs=18,
    validation_data=val_data,
    callbacks=[loss_logger, checkpoint_callback, early_stopping]
)

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