采用 GitHub 上 star 比较高的一个开源实现 https://github.com/yhenon/pytorch-retinanet
在 anaconda 中新建了一个环境, 因为一开始并没有新建环境, 在原有的环境里, 遇到了 pytorch,numpy 等版本问题, 尤其是一开始用 0.1.2 版的 pytorch, 有一个TH路径不对, 导致编译错误, 这是我最讨厌的错误, 遇到编译错误往往一脸懵逼, 如果NMS部分不用编译, 直接用 python 实现就好了, 当然那样速度可能会慢很多.
先记录下我的各个包的版本
- cffi 1.12.2
- cudatoolkit 9.0
- cudnn 7.3.1
- Cython 0.29.7
- matplotlib 3.0.3
- numpy 1.15.4
- pytorch 0.4.0
- torchvision 0.2.1
当然了, 其他版本也可以, 但是这个版本一定是可行的.
然后准备训练 coco, 首先需要下载 coco, 这里采用 wget 下载, 几个压缩包的地址链接 https://blog.csdn.net/daniaokuye/article/details/78699138#commentsedit.
采用 wget 下载意外地很慢, 于是采用迅雷, 意外地很快
网络开始训练, 没有采用预训练权重, 既没有用 coco 训练好的, 也没有用 resnet 的预训练权重 (下载起来太慢了)
python train.py --dataset coco --coco_path ../coco --depth 50
训练截图
采用2个图片作为一个 batch 训练,GPU占用
batchsize 为2, 训练一个 epoch 大约6个小时, 按照代码中默认的100个 epoch, 恐怕得600个小时, 一个月了
幸好源代码中提供了训练好的 coco 权重, 可以为我们所用, 那就先看一下训练好的效果, 调用可视化代码
python visualize.py --dataset coco --coco_path ../coco --model ./coco_resnet_50_map_0_335_state_dict.pt
效果
检测效果还是 ok 的~
这个版本的实现代码量在 2000 行左右, 很适宜阅读, 尤其是与 Mask R-CNN(matterport 版, 大概 6000 行) 相比
可视化代码稍加改造, 就可以作为一个目标检测器使用了, 棒!
来源: https://www.cnblogs.com/wzyuan/p/10822162.html