保存 YOLOv3 模型状态和加载

Question

I am trying to train YOLOv3 model following ImageAi documentation with my custom dataset.But the training is taking too much time that Google Colab time is not enough for this. Now how can I save Model state and load it after 50/60 epoch completion? As I am a beginner I am not getting Tensorflow Model Checkpoint t

Here is a sample of the code:

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()

Answer 1

You should use a ModelChekpoint as callback of your method. But you use a custom class that already handle this.

If you have a look at the code on there github you can see that they use some custom callbacks

They save the model after each epoch by only the best one (not the last one). The model should be saved in the folder <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-")

If you want to save the last iteration you have to overwrite the method for custom callbacks and change save_best_only to False .