PyTorch版EfficientDetをROCm環境で実行してみた

PyTorch版EfficientDet1のカメラで撮影した画像に推論結果を重畳してくれるデモをROCm環境で動かしたので手順などをメモ。

環境

dockerイメージ

Docker Hubでrocm/pytorchの中からそれっぽいのを選択。 Ubuntu18.04はROCm2.10が一番新しかった。

docker pull rocm/pytorch:rocm2.10_ubuntu18.04_py3.6_pytorch_profiling

ホスト側の準備[2020/05/12 追記]

以下のエラーが出てしまうのでxhostコマンドで許可を出しておく。

No protocol specified
: cannot connect to X server :0

ホスト側(コンテナの外)で実行する。

xhost local:

コンテナの起動コマンド

  • ホームディレクトリを/dataにマウント
  • X11アプリを動かすため/tmp/.X11-unixをコンテナに見せる
  • カメラをデモで使うので/dev/video0をコンテナに見せる
  • cv2.imshow()、cv2.waitKey()を動かすためDISPLAY変数、QT_X11_NO_MITSHMを設定する
  • 他の設定はコンテナ内でGPUを使うため
docker run -it --name pytroch_rocm -v ~/:/data -v /tmp/.X11-unix/:/tmp/.X11-unix/ --privileged --device /dev/video0 --device=/dev/kfd --device=/dev/dri --group-add video -e DISPLAY=$DISPLAY -e QT_X11_NO_MITSHM=1 rocm/pytorch:rocm2.10_ubuntu18.04_py3.6_pytorch_profiling

※1:コンテナ起動時にカメラは接続済みの状態にしておく
※2:QT_X11_NO_MITSHMが無いとcv2.waitKey()で以下のエラーになる

X Error: BadAccess (attempt to access private resource denied) 10
  Extension:    130 (MIT-SHM)
  Minor opcode: 1 (X_ShmAttach)
  Resource id:  0x4eb

Docker上で追加作業

X11の準備

デモの中でOpenCVのウィンドウを表示するのでx11-appsを追加した。

apt-get update
apt-get install x11-apps

Pythonパッケージ

pip installでエラーになるケースがあるので、いつもpipとsetuptoolsを最新にしている。システムのpip3を更新すると環境を壊してしまうことがあるのでvenv環境を作ってそちらに入れている。

apt-get install python3-venv
python3 -m venv --system-site-packages env_rocm
source env_rocm/bin/activate
pip install --upgrade pip
pip install --upgrade setuptools

EfficientDetの依存パッケージ

デモでOpenCVを使っているのと、COCO APIを使っているので一緒にインストールする。

pip install -r requirements.txt
pip install opencv-python
pip install git+https://github.com/cocodataset/cocoapi.git#subdirectory=PythonAPI

EfficientNetの重みデータをキャッシュディレクトリに置く

EfficientDetのGitHubサイトにはEfficintDet部分の重みデータへのリンクが貼られているが、どうやらバックボーンネットワークのEfficientNetの重みを別途実行時にダウンロードするらしい。

現状ソースコードに埋め込まれているダウンロードURLがリンク切れしているようなので、EfficientNetのPyTorch実装GitHub2サイトからダウンロードして、キャッシュディレクトリに置く。

# wgetなどでダウンロードしておく

mkdir -p /root/.cache/torch/checkpoints
cd /root/.cache/torch/checkpoints
ln -s /data/python3.6/efficientnet-b0-355c32eb.pth efficientnet-b0-355c32eb.pth

ソースコード修正

dockerイメージに含まれているtorchvisionはCUDA実装がビルドされていないようでNMS呼び出しで以下のエラーが出る。

  File "demo.py", line 177, in <module>
    detect.camera()
  File "demo.py", line 155, in camera
    show_image = self.process(img=img)
  File "demo.py", line 81, in process
    scores, classification, transformed_anchors = self.model(img)
  File "/root/.local/lib/python3.6/site-packages/torch/nn/modules/module.py", line 539, in __call__
    result = self.forward(*input, **kwargs)
  File "/data/python3.6/EfficientDet.Pytorch/models/efficientdet.py", line 83, in forward
    transformed_anchors[0, :, :], scores[0, :, 0], iou_threshold=self.iou_threshold)
  File "/root/.local/lib/python3.6/site-packages/torchvision/ops/boxes.py", line 31, in nms
    return torch.ops.torchvision.nms(boxes, scores, iou_threshold)
RuntimeError: Not compiled with GPU support (nms at /tmp/pip-wnkd089r-build/torchvision/csrc/nms.h:20)

torchvisionのHIPIFYのやり方がよくわかってないのでビルドは諦めて呼び出し側でtorch.tensorをいったんCPUに戻すようにしたらエラー回避できた。

diff --git a/models/efficientdet.py b/models/efficientdet.py
index 43357c9..83d60d9 100644
--- a/models/efficientdet.py
+++ b/models/efficientdet.py
@@ -80,7 +80,7 @@ class EfficientDet(nn.Module):
             transformed_anchors = transformed_anchors[:, scores_over_thresh, :]
             scores = scores[:, scores_over_thresh, :]
             anchors_nms_idx = nms(
-                transformed_anchors[0, :, :], scores[0, :, 0], iou_threshold=se
lf.iou_threshold)
+                transformed_anchors[0, :, :].cpu(), scores[0, :, 0].cpu(), iou_threshold=self.iou_threshold).cuda()
             nms_scores, nms_class = classification[0, anchors_nms_idx, :].max(
                 dim=1)
             return [nms_scores, nms_class, transformed_anchors[0, anchors_nms_idx, :]]

実行

実行は作者様のGitHubサイトの通りに実行。EfficientDetの重みデータは、あらかじめGitHubに貼られているリンクからダウンロードしておく。

python demo.py --threshold 0.5 --iou_threshold 0.5 --cam --score --weight ../checkpoint_VOC_efficientdet-d0_268.pth

  1. A Pytorch Implementation of EfficientDet Object Detection, https://github.com/toandaominh1997/EfficientDet.Pytorch

  2. EfficientNet PyTorch, https://github.com/lukemelas/EfficientNet-PyTorch/releases