保存 YOLOv3 模型状态和加载答案

【问题标题】：Saving YOLOv3 Model State and Loading保存 YOLOv3 模型状态和加载
【发布时间】：2021-01-06 00:44:49
【问题描述】：

我正在尝试使用我的自定义数据集按照 ImageAi 文档训练 YOLOv3 模型。但是训练花费了太多时间，Google Colab 的时间不足以完成此任务。现在如何保存模型状态并在 50/60 epoch 完成后加载它？因为我是初学者，所以我没有获得 Tensorflow 模型检查点 t

以下是代码示例：

from imageai.Detection.Custom import DetectionModelTrainer 

trainer = DetectionModelTrainer() 
trainer.setModelTypeAsYOLOv3() trainer.setDataDirectory(data_directory="/content/drive/My Drive/Dataset") 
trainer.setTrainConfig(object_names_array=["obj1","obj2"], batch_size=4, num_experiments=421) 
trainer.trainModel()

【问题讨论】：

您应该放一个代码示例并说明您使用的是哪个版本的 tensorflow
from imageai.Detection.Custom import DetectionModelTrainer trainer = DetectionModelTrainer() trainer.setModelTypeAsYOLOv3() trainer.setDataDirectory(data_directory="/content/drive/My Drive/Dataset") trainer.setTrainConfig(object_names_array= ["obj1","obj2"], batch_size=4, num_experiments=421) trainer.trainModel()
!pip3 install tensorflow==1.13.1

标签： python tensorflow detection imageai yolov5

【解决方案1】：

您应该使用 ModelChekpoint 作为方法的回调。但是您使用的自定义类已经处理了这个问题。

如果您查看 github 上的代码，您会发现他们使用了一些 custom callbacks

他们在每个 epoch 之后只保存最好的模型（不是最后一个）。模型应该保存在文件夹/models/detection_model-:

self.__model_directory = os.path.join(data_directory, "models")
self.__train_weights_name = os.path.join(self.__model_directory, "detection_model-")

如果您想保存最后一次迭代，您必须覆盖自定义回调的方法并将save_best_only 更改为False。

【讨论】：