ssd.pytorchをpytorch 1.5.1で動くように修正した

ssdの元論文と同等の精度のpytorch実装にamdegroot/ssd.pytorchがあるのですが、pytorchの古いバージョンでしか動きません。CUDAのversionも合わせないとGPUが使えないので古いバージョンで動かすのは大変です(というか僕は出来ませんでした)。なのでpytorch 1.5.1で動くように修正しました。変更内容はここで確認出来ます。試してみる場合はhttps://github.com/nnabeyang/ssd.pytorchのブランチpytorch1.5.1を取ってくるのが一番手っ取り早いと思います。

この状態で学習して出来たモデルのmAPは0.77で、古いバージョンと同等の性能でした。

AP for aeroplane = 0.7946
AP for bicycle = 0.8467
AP for bird = 0.7603
AP for boat = 0.7286
AP for bottle = 0.4905
AP for bus = 0.8546
AP for car = 0.8628
AP for cat = 0.8865
AP for chair = 0.6305
AP for cow = 0.8270
AP for diningtable = 0.7645
AP for dog = 0.8494
AP for horse = 0.8707
AP for motorbike = 0.8396
AP for person = 0.7927
AP for pottedplant = 0.4953
AP for sheep = 0.7737
AP for sofa = 0.8142
AP for train = 0.8523
AP for tvmonitor = 0.7710
Mean AP = 0.7753

今回の修正のほとんどは参考に載せているページに載っている https://github.com/ponta256/pytorch-ssd7から取ってきています。学習の仕方やmAPの計算方法もそこのページに書かれています。

環境

• Ubuntu 20.04
• Python 3.6.9
• Pytorch 1.5.1
• CUDA 10.2

参考

• Pytorch版SSDを改造してベースネットを入れ替えてみる|株式会社Sigfoss