複数GPU環境におけるPytorchでの使用GPU指定に関するTips

cuda環境でコードを走らせようとするとき、モデルやデータをGPUに載せる操作を以下のように記述していることがある。

model = model.cuda()
data = data.cuda()
batch = batch.cuda()

Google Colaboratoryなど使用できるGPUがそもそも1つしかない場合はこのままで何ら問題はないが、複数GPUを使用できる環境であったり、他人と分担してGPUを使わなければならなかった場合はこのままだと困る場合がある。

使用GPUを指定できないと困るケース(例)

• 複数人で共同で使用している環境である場合

例えばcuda:0とcuda:1がある環境で同時に複数の人がcuda:0にアクセスしてしまうとCUDA OUT OF MEMORYのエラーが出てきてしまったり、動作が想定よりも遅くなってしまう場合がある。

• 使用すべきGPUが指定されている場合

そのため、研究室などで使用可能なGPUが何台かあって共同利用しなくてはならないときに「○○さんは何番と何番使ってください」と言われる場合がある。しかし、先述のような.cuda()を用いてGPUデバイスに送っていると常にデフォルトのcuda:0が選択されてしまうので注意が必要

以上の話は、機械学習のための全てのコードを自分で書いていたら、GPUに送る必要のあるオブジェクトすべてでcuda:n(nは使用するGPUの番号)を指定すれば済む話である。

GPU_ID = 4 #使用するGPUのID
device = torch.device(f'cuda:{GPU_ID}' if torch.cuda.is_available() else 'cpu')
model = model.to(device)
data = data.to(device)
batch = batch.to(device)

しかし、使用しているコードの一部がオープンソースなどの外部ライブラリに依存していて外部ライブラリが.cuda()を使用して転送していた場合、それら全部書き直すのは面倒だしそもそもできない場合もある。（筆者は今回このケースにぶち当たった。）

そこで、.cuda()をすべて.to(device)に書き換えなくても何とかなる方法をここでまとめる。

使用GPUを指定する方法

• 環境変数CUDA_VISIBLE_DEVICESで指定する

指定方法はスクリプト実行時に指定する方法とスクリプト内で指定する方法がある。
複数のGPUを指定する場合はカンマ(,)で区切る。

$ export CUDA_VISIBLE_DEVICES=4 #1つのGPUを指定する場合
$ export CUDA_VISIBLE_DEVICES=4,5,6,7 #複数のGPUを指定する場合

import os

os.environ['CUDA_VISIBLE_DEVICES'] = '4' #1つのGPUを指定する場合
os.environ['CUDA_VISIBLE_DEVICES'] = '4,5,6,7' #複数のGPUを指定する場合

ただし、ここで複数のGPUを指定していた場合に、そのうちの特定のGPU(例えば4番)を使いたい場合は、cuda:4ではなくcuda:0としなければならない点に注意が必要。

恐らく番号順にインデックス0始まりで指定できるようになっていると考えられる。この例だとcuda:1 -> 5番のGPU,cuda:2 -> 6番のGPU,...というように。cuda:4を指定するとlistでいうところのout of rangeのGPUを指定してしまっていることになる(4つのGPUしか指定していないのでインデックスは0~3)ので、invalid device ordinalというエラーが出力される。

• torch.cuda.set_device(hogehoge)を使う

筆者のケースではなぜか上述の方法だと想定していないGPUにもデータが乗ってしまう場合があったので、これで解決した。

torch.cuda.set_device(5) #使用したいGPUの番号を入れる

これをスクリプトの始めの方に書いておくと、指定したGPUがいわゆるデフォルトのデバイスになるので、.cuda()を使ったとしてもちゃんと指定したGPUにオブジェクトが乗ってくれた。ありがたい。

参考:https://pytorch.org/docs/stable/generated/torch.cuda.set_device.html