保存YOLOv3模型状态并加载

问题描述

我正在尝试使用自定义数据集按照ImageAi文档训练YOlov3模型，但是训练时间太长，以至于Google Colab时间还不够。现在，如何保存模型状态并在50/60纪元完成后加载它？因为我是初学者，所以我没有得到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()

解决方法

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

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

在每个时期之后，它们仅以最好的一个（而不是最后一个）保存模型。模型应保存在 / 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。