MJUN Tech Note

初めてのissueとPull ReqでPyTorchとYOLOv5に貢献した話

こんにちは.今回は私が初めて立てたissueがPyTorchのバグを修正に繋がり, 新しいCI/CDテスト項目を増やすに至った話と, このバグをもとに初めてのPull RequestをYOLOv5に送ってcontributerになった話を 書きたいと思います.

個人的な話になりますが,大きなOSSに貢献するのは初めてなので, 記録に残して,OSSに貢献する流れが伝わればと思い書くことにします.

事の経緯

普段,深層学習手法の実装にPyTorchを利用しているのですが, 使用するGPUの切り替えを以下のようなコードで行っていました.

import os
import torch

os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "1"

# print using GPU Info
print(f"Using GPU is CUDA:{os.environ['CUDA_VISIBLE_DEVICES']}")

for i in range(torch.cuda.device_count()):
    info = torch.cuda.get_device_properties(i)
    print(f"CUDA:{i} {info.name}, {info.total_memory / 1024 ** 2}MB")

device = torch.device("cuda:0")

上記のコードではPyTorchでGPU割り当てに, 環境変数CUDA_VISIBLE_DEVICESを使用しています. こうすることで,実行プロセス内で指定したGPU以外をマスクすることができるので, 安全にリソースを扱えます.(例えばコード上でどこでもcuda:0を書いても安全)

PyTorchではtorch.cudaモジュールが呼び出されるまではCUDA Deviceの初期化は 行われないので,環境変数をコード上で変更したあとにtouch.cudaを呼び出せば, 環境変数が適用されます.(いわゆるlazy load)

環境変数なので,実行時にもCUDA_VISIBLE_DEVICES=0 python train.pyというように指定できます.

このコードがPyTorch Version 1.11.0までは動いていたのですが, 1.12.0から動かなくなってしまいました.

具体的な例をあげます. 以下のような計算機があったとします.

CUDA:0 NVIDIA RTX A6000, 48685.3125MB
CUDA:1 NVIDIA GeForce RTX 3090, 24268.3125MB

この計算機でCUDA:1のRTX 3090を使いたいと考え, 上記のコードをVer 1.12.0と1.11.0で動かします. すると出力は以下のようになります.

  • Ver 1.11
Using GPU is CUDA:1
CUDA:0 NVIDIA GeForce RTX 3090, 24268.3125MB
  • Ver 1.12
Using GPU is CUDA:1
CUDA:0 NVIDIA RTX A6000, 48685.3125MB
CUDA:1 NVIDIA GeForce RTX 3090, 24268.3125MB

上記の通り,Ver 1.12ではコード上で書き換えたCUDA_VISIBLE_DEVICESの環境変数が 適用されていないことがわかります.

私はこの方法でGPUの割当を行っていたのでとても困りました.

YOLOv5のIssue

物体検出手法であるYOLOv5は PyTorchで実装されています.YOLOv5も同じような方法でGPUの割当を行っていたのを 思い出したので,YOLOv5で同じような問題が起きていないかIssueを探すことにしました.

すると,リポジトリのAuthorであるGlenn Jocher氏が以下のようなIssueを 上げていました.

YOLOv5 issues with torch==1.12 on Multi-GPU systems

まさしく,私と同じ問題を抱えていました. そこで,問題を再現可能な小さいコードを示しました.

import os
import torch

os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "1"

# print using GPU Info
print(f"Using GPU is CUDA:{os.environ['CUDA_VISIBLE_DEVICES']}")

for i in range(torch.cuda.device_count()):
    info = torch.cuda.get_device_properties(i)
    print(f"CUDA:{i} {info.name}, {info.total_memory / 1024 ** 2}MB")

そして,以下のようにimport torchの前に環境変数を指定すれば 上手く動作することもわかりました.

import os

os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "1"

import torch

その後もimportlibを使ったモジュールのreload等を試しましたが問題は解決せず, これはPyTorch自体の問題であると考え,PyTorchに人生初のIssueを投稿することにしました.

PyTorchのIssue

私がPyTorchのリポジトリに投稿したIssueは以下です.

[1.12] os.environ[“CUDA_VISIBLE_DEVICES”] has no effect

投稿したところ,この問題がhigh priorityに指定されたため,とても緊張した覚えがあります.

PyTorchのMain Contributerとやり取りする中で,この問題を再現できる人とできない人がいたため, 再現しやすいようにDockerfileの例を出しながら,どのような場合に問題が起きるかを調査しました.

最終的に,この問題を発生させる最小コードは以下のようになりました.

import os
import torch

os.environ["CUDA_VISIBLE_DEVICES"] = "32"
print(torch.__version__, torch.cuda.device_count())
# ver 1.11
$ python cudev.py 
1.11.0+cu113 0

# ver 1.12
$ python3 cudev.py
1.12.0+cu113 7

CUDA_VISIBLE_DEVICESを32というありえない数字にすることで, GPUの割当を0にする試みです.このコードも1.12では期待した動作をしませんでした.

すると,Contributorであるmalfet氏が 問題の原因を発見しました.

ver 1.11から1.12までのcommitで,import torchでモジュールを読み込むまでの間に torch.cudaを呼び出している箇所があり,これが原因でCUDA Deviceの初期化が意図せず 行われていることがわかりました.一度CUDA Deviceの初期化が行われると 再度初期化を行うことはできないので,torch.cudaを呼び出さないように修正されました.

そして,この問題が再発しないように, PyTorchのテスト項目として,以下のコードが追記されました.

@unittest.skipIf(TEST_WITH_ROCM, "ROCm doesn't support CUDA_VISIBLE_DEVICES")
@unittest.skipIf(TEST_MULTIGPU, "Testing on one GPU is sufficient")
def test_lazy_init(self):
    """ Validate that no CUDA calls are made during `import torch` call"""
    from subprocess import check_output
    test_script = "import os; import torch;os.environ['CUDA_VISIBLE_DEVICES']='32';print(torch.cuda.device_count())"
    rc = check_output([sys.executable, '-c', test_script]).decode("ascii").strip()
    self.assertEqual(rc, "0")

この修正はPyTorch Version 1.12.1で適用される予定です.

YOLOv5にPull Request

PyTorchに投稿したIssueを踏まえて, この問題をYOLOv5のコード内で解決することは難しそうでしたが, PyTorchの次期バージョンで修正されることがわかりました.

そのため,YOLOv5のリポジトリにあるrequirements.txtに以下の変更を加えて 修正するPull Requestを贈りました.

- torch>=1.7.0
- torchvision>=0.8.1

+ torch>=1.7.0,!=1.12.0  # https://github.com/ultralytics/yolov5/issues/8395
+ torchvision>=0.8.1,!=0.13.0 # https://github.com/ultralytics/yolov5/issues/8395

この問題は特定のバージョンのPyTorchで起こり,将来の1.12.1で修正されるので, Ver 1.12と対応するtorchvisionのバージョンを除外することで問題を解決することにしました.

このPull Requestは無事mergeされ,私は初めてOSSにPull Requestを送ってmergeされました.

まとめ

初めてのIssue, Pull Requestだったため,先人のサイトを読みながら慎重に やり取りを進めました.

PyTorchのContributorがreproと言って気にしていたように, Issue報告をする際は問題を再現する最小のコードを提供することが大切であると 痛感しました.これは普段の開発でも大切なことです.

下世話な話ですが,先にPyTorchの問題箇所を発見できていれば PyTorchのContributorになれたのに…なんて思ったりしました.

以上,PyTorchとYOLOv5に貢献した話でした.