SSD-Pytorch
作者爲pytorch中文網維護者之一,網站會發許多深度學習相關文章,歡迎pytorch使用者多多支持.
項目Github地址:Github 歡迎star, fork
一. 項目簡介
-
項目是在實際項目上整理精簡而來,主要針對單機多卡環境實現.
-
項目整體結構清晰明瞭,對訓練過程在不失靈活性的前提下進行了包裝.訓練,測試,檢測過程均只需簡單幾行即可實現.
-
如需深入瞭解,可對照項目文件結構進行理解,項目整體均有詳細中文備註.
-
完整項目地址:Github,歡迎star, fork.
-
github連接較慢的,可以去Gitee(國內的代碼託管網站),也有完整項目.
1.1 項目特點
1.結構化實現.模型各部件均分離實現,模型結構更爲清晰明瞭,便於理解以及後續修改使用.
2.針對單機多卡環境實現.簡單直觀的多GPU調用實現. 只需通過列表指明需調用的GPU序號即可實現多GPU環境下的訓練與測試.如: train_devices=[0,1,3]即可調用第1,2,4三塊GPU用於訓練.
3.全中文註釋.每部分均有詳細的中文註釋,是你學習SSD,加深理解的不二之選.
1.2 Requirements
1.pytorch
2.opencv-python
3.torchvision >= 0.3.0
4.Vizer (在代碼實現上也借鑑了lufficc大神的一些思想)
5.visdom
1.3 項目文件結構
| 文件夾 | 文件 | 說明 |
|---|---|---|
| Data | 數據相關 | |
| Dataloader | 數據加載器類’Our_Dataloader’, ‘Our_Dataloader_test’ | |
| Dataset_VOC | VOC格式數據集類 | |
| Transfroms | 數據Transfroms | |
| Transfroms_tuils | Transfroms子方法 | |
| Model | 模型相關 | |
| base_models/VGG | 現只有vgg模型,後期會更新其他模型 | |
| structs/Anchors | 默認檢測框生成器 | |
| structs/MutiBoxLoss | SSD損失函數 | |
| structs/PostProcess | 後處理 | |
| structs/Predictor | 分類及迴歸網絡 | |
| evaler | 驗證器,用於在數據集上對模型進行驗證(測試),計算ap,map | |
| ssd_model | SSD模型類 | |
| trainer | 訓練器,用於在數據集上訓練模型 | |
| Utils | 各種工具 | |
| boxs_op | 各種框體操作,編碼解碼,IOU計算,框體格式轉換等 | |
| Weights | 模型權重存放處 | |
| pretrained | 預訓練模型權重存放處 | |
| trained | 訓練過程中默認模型存放處 | |
| ---- | Configs.py | 配置文件,包含了模型定義,數據以及訓練過程,測試過程等的全部參數,建議備份一份再進行修改 |
| ---- | Demo_train.py | 模型訓練的例子,訓練過程中的模型會保存在Weights/Our/ |
| ---- | Demo_eval.py | 模型測試的例子,計算模型ap,map |
| ---- | Demo_detect_one_image.py | 檢測單張圖片例子 |
| ---- | Demo_detect_video.py | 視頻檢測例子,傳入一個視頻,進行檢測 |
二. Demo例子
本項目Demo均是完整的訓練,測試,檢測過程,可直接運行.
2.1 訓練train
本項目採用visdom進行訓練過程可視化
visdom安裝 及運行:
visdom安裝
pip install visdom 安裝
python -m visdom.server 運行
打開網頁`127.0.0.1:8097` 進行訓練過程可視化
(如未運行visdom,會報警告,但是不影響訓練過程.建議通過visdom進行可視化)
訓練過程通過訓練器進行配置,訓練器支持單機單卡,單機多卡訓練;不支持多機多卡以及cpu.
# 實例化模型. 模型的具體各種參數在Config文件中進行配置
net = SSD(cfg)
# 將模型移動到gpu上,cfg.DEVICE.MAINDEVICE定義了模型所使用的主GPU
# 模型的參數更新會在主GPU上進行.
net.to(cfg.DEVICE.MAINDEVICE)
# 初始化訓練器,訓練器參數已通過cfg進行配置;也可傳入參數進行配置,但不建議
trainer = Trainer(cfg, max_iter=None, batch_size=None,
train_devices=None, model_save_step=None,
model_save_root=None, vis = None, vis_step=None)
# 訓練器開始在 數據集上訓練模型
trainer(net, train_dataset)
2.2 測試Eval
對模型進行測試,測試器支持單機單卡,單機多卡測試,但不支持多機多卡,cpu.
net = SSD(cfg)
# 將模型移動到gpu上,cfg.DEVICE.MAINDEVICE定義了模型所使用的主GPU
net.to(cfg.DEVICE.MAINDEVICE)
# 模型從權重文件中加載權重
net.load_pretrained_weight('Weights/pretrained/vgg_ssd300_voc0712.pkl')
# 初始化驗證器,驗證器參數通過cfg進行配置;也可傳入參數進行配置,但不建議
evaler = Evaler(cfg, eval_devices=None)
# 驗證器開始在數據集上驗證模型
ap, map = evaler(model=net, test_dataset=test_dataset)
print(ap)
print(map)
2.3 檢測Detect
檢測過程本項目給出倆個例子,分別爲單張圖片檢測與視頻檢測
2.3.1 圖片檢測
單張圖片檢測,通過調用模型Detect_single_img方法實現.支持gpu與cpu,依模型運行環境而定.
# 實例化模型
net = SSD(cfg)
# 使用cpu或gpu
net.to('cuda')
# 模型從權重文件中加載權重
net.load_pretrained_weight('Weights/pretrained/vgg_ssd300_voc0712.pkl')
# 打開圖片
image = Image.open("Images/000133.jpg")
# 進行檢測, 分別返回 繪製了檢測框的圖片數據/迴歸框/標籤/分數.
drawn_image, boxes, labels, scores = net.Detect_single_img(image=image,score_threshold=0.5)
plt.imshow(drawn_image)
plt.show()
2.3.2 視頻檢測
# 實例化模型
net = SSD(cfg)
# 使用cpu或gpu
net.to('cuda')
# 模型從權重文件中加載權重
net.load_pretrained_weight('Weights/pretrained/vgg_ssd300_voc0712.pkl')
video_path = 'aaa.mp4' # 現在忙於其他事,項目後期會上傳視頻檢測樣例
# 進行檢測,
# save_video_path=None,則不保存視頻; save_video_path=bbb.mp4 將檢測後視頻保存爲bbb.mp4
# show=True,實時顯示檢測結果
net.Detect_video(video_path=video_path, score_threshold=0.02, save_video_path=None, show=True)
說明
本項目是在實際項目上整理精簡而來,現配置文件爲在voc2007上進行訓練,如有其他需求,請自行修改.
完整項目地址:Github,歡迎star, fork.
僅限用於交流學習,如需引用,請聯繫作者.