Win10下ImageAI-gpu訓練自己的數據集

1.先按《Win10下ImageAI-gpu(tensorflow-gpu)環境安裝(cuda及cuDNN自動安裝)及實時對象檢測》安裝好。

2.下載pretrained-yolov3.h5,放到temp.py同一目錄下。

3.下載hololens數據集,放到temp.py同一目錄下。

4.編寫代碼開始訓練。

from imageai.Detection.Custom import DetectionModelTrainer
trainer = DetectionModelTrainer()
trainer.setModelTypeAsYOLOv3()
trainer.setDataDirectory(data_directory="hololens")
trainer.setTrainConfig(object_names_array=["hololens"], batch_size=2, num_experiments=100, train_from_pretrained_model="pretrained-yolov3.h5")
trainer.trainModel()

batch_size=2是關鍵,否則可能出現以下錯誤

Resource exhausted: OOM when allocating tensor with shape[8,1024,14,14] and type float on /job:localhost/replica:0/task:0/device:GPU:0 by allocator GPU_0_bfc

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章