0

我正在尝试使用我的自定义数据集按照 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()
4

1 回答 1

0

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

如果您查看那里 github 上的代码,您会发现它们使用了一些自定义回调

他们在每个 epoch 之后只保存最好的模型(不是最后一个)。模型应保存在文件夹 <data_directory>/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_onlyFalse.

于 2020-09-21T09:34:39.493 回答