M1 MacのGPUでYOLOXを動かしてみた

公開日：2022/09/11

どうやらPyTorchがM1 Mac (Apple silicon) 対応したということで、インストールして試してみました。

PyTorchのインストール

まずはインストールですが、PyTorch公式サイトのGet StartedからNightly -> Mac -> Pip -> Python を選択するとpipでのインストールコマンドが出てきます。

私の場合は以前にTensorflow環境を作った際にvenvを導入していましたので、PyTorch用に新たなvenv環境を作ってインストールを進めました。

# venv作成
python3 -m venv pytorch

# venvをactivate
source pytorch/bin/activate

# pipが古そうだったのでアップデート
pip3 install -U pip

# PyTorchをインストール
pip3 install --pre torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/nightly/cpu

# Python3を実行し、M1 MacのGPUモードであるmpsに切り替えができるかを確認
python3

>>> import torch
>>> torch.device('mps')

ここまでの作業で、かなりあっさりインストールが完了しました。

なんでそんなターミナル画面なの？というと、自宅にPCがありすぎて今回はWindowsからSSHとVNCでMacに入って使っているからです。
普通にMacのターミナルから操作できます。

YOLOXをGPU対応で動作させるまで

PyTorchのインストールが完了したので、何かで試してみようと思い、YOLOXで確認することにしました。
git cloneして一式をダウンロードしたら、まずは必要なライブラリ等を入れていきます。
ここは本題から外れるため割愛します。

修正する前に、CPUでの実行がどの程度なのか試してみます。
モデルはyolox_s.pthを使います。

フレームあたり0.13～0.14秒ぐらい。
Core i7より少し遅いぐらいでしょうか。

そしてここからGPU対応のために手を加えていきます。
基本的にはcudaを使っている部分をmpsに変えていきます。

tools/demo.pyで、cudaの設定をしている部分を探します。

if self.device == "gpu":
   (略)
elif self.device == "mps":
    img = img.to(torch.device("mps"))

もう1か所(元ファイルの行番号なので上記コード追加前の行番号です)

if args.device == "gpu":
   (略)
elif args.device == "mps":
    model.to(torch.device("mps"))

早速実行してみると、何やらエラーが…

ValueError: invalid type: 'torch.mps.FloatTensor'

該当箇所はyolox/models/yolo_head.pyのLine246

grids = torch.cat(grids, dim=1).type(dtype)

調べてみるとtorch.catはmpsに対応していないらしいということがわかったので、少々強引ではあるもののここのブロックを以下のように書き換え

if dtype.startswith("torch.mps"):
    grids = torch.cat(grids, dim=1).to("cpu")
    strides = torch.cat(strides, dim=1).to("cpu")
    outputs = outputs.to("cpu")
else:
    grids = torch.cat(grids, dim=1).type(dtype)
    strides = torch.cat(strides, dim=1).type(dtype)

すると、あっさり通りました。(もっと苦労するかと思ってた)

0.035秒前後ならそこそこ実用的なレベルで速いのではないでしょうか。

実行結果も問題なくこのとおり

ちなみに今回使用した動画素材は、こちらを30fpsにリフレームしたものです。
今回作ったコードはYOLOXのリポジトリをフォークしてgithubのj-ohashi/YOLOX-M1-Macに上げています。

もしかするとmpsでもcudaの実装にあるようなempty_cacheとか必要なのかもしれませんが、現時点ではとりあえず触ってみたレベルであまり細かいところまで調べ切れていません。
また今回の実装には使っていませんが、cudaの使用可否を確認するtorch.cuda.is_available()に相当するのは、torch.backends.mps.is_available()となります。
個人的に忘れそうなのでここに書いておきます。

まとめとしては、昨今、NVIDIA製GPUが入手しにくかったりJetsonやラズパイもかなり高騰している中、今のタイミングであればM1 Mac miniであれば10万円未満で比較的容易に手に入れることができます。
もしかするとエッジデバイスとしてMac miniを使うのも有りなのでは？と個人的には期待しています。

M1 MacのGPUでYOLOXを動かしてみた

PyTorchのインストール

YOLOXをGPU対応で動作させるまで

M1 MacのGPUでYOLOXを動かしてみたにコメントする コメントをキャンセル

M1 MacのGPUでYOLOXを動かしてみたにコメントするコメントをキャンセル