Pytoorch輕松學 – RetinaNet自定義對象檢測

RetinaNet網(wǎng)絡介紹

2017發(fā)布，實現(xiàn)了一階段網(wǎng)絡首次在精度方面超過二階段網(wǎng)絡的經(jīng)典網(wǎng)絡，作者最大的一個創(chuàng)新就是在訓練損失函數(shù)方面，論文比較了CE、BCE、以及論文提出感知損失函數(shù)（FL），最后說明感知損失可以有效解決一階段網(wǎng)絡訓練中的樣本不平衡現(xiàn)象，從而取得更佳的訓練效果。論文中提出的感知損失函數(shù)如下：

最終RetinaNet網(wǎng)絡結(jié)構如下：

數(shù)據(jù)集準備與制作

自己百度收集了一個無人機與飛鳥的數(shù)據(jù)集，其中訓練集270張圖像，測試集26張圖像。

使用labelImg工具完成標注，工具下載地址：

https://gitee.com/opencv_ai/opencv_tutorial_data/tree/master/tools

三：模型訓練

制作好數(shù)據(jù)集之后，模型訓練就成為一件很簡單事情，基于OpenMV工具軟件，零代碼即可實現(xiàn)模型訓練。運行下面界面如下：

總計訓練了25個輪次以后，發(fā)現(xiàn)效果已經(jīng)是相當?shù)牟诲e了，直接導出ONNX格式RetinaNet模型文件。把模型轉(zhuǎn)換為ONNX格式，Pytorch是原生支持的，只需要把通過torch.onnx.export接口，填上相關的參數(shù)，然后直接運行就可以生成ONNX模型文件。相關的轉(zhuǎn)換代碼如下：

model=tv.models.detection.retinanet_resnet50_fpn(pretrained=True)
dummy_input=torch.randn(1,3,1333,800)
model.eval()
model(dummy_input)
im=torch.zeros(1,3,1333,800).to("cpu")
torch.onnx.export(model,im,
"retinanet_resnet50_fpn.onnx",
verbose=False,
opset_version=11,
training=torch.onnx.TrainingMode.EVAL,
do_constant_folding=True,
input_names=['input'],
output_names=['output'],
dynamic_axes={'input':{0:'batch',2:'height',3:'width'}}
)

運行時候控制臺會有一系列的警告輸出，但是絕對不影響模型轉(zhuǎn)換，影響不影響精度我還沒做個仔細的對比。模型轉(zhuǎn)換之后，可以直接查看模型的輸入與輸出結(jié)構，圖示如下：

推理運行

推理部分的代碼很簡單，只有三十幾行，Python就是方便使用，這里最需要注意的是輸入圖像的預處理必須是RGB格式，需要歸一化到0~1之間。對得到的三個輸出層分別解析，就可以獲取到坐標（boxes里面包含的實際坐標，無需轉(zhuǎn)換）。基于OpenMV工具軟件，可以實現(xiàn)一鍵零代碼推理演示，效果如下：