我正在尝试使用自定义数据集按照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()
答案 0 :(得分:0)
您应该使用ModelChekpoint作为方法的回调。但是,您可以使用已经处理此问题的自定义类。
如果您查看github上的代码,您会发现他们使用了custom callbacks
在每个时期之后,它们仅以最好的一个(而不是最后一个)保存模型。模型应保存在
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
。