Pytorch LightnigでGPUを指定する方法

こんにちは、九州大学修士２年の浅海です。
Pytorch lightningで指定のGPUを使用する設定に、つまづいたのでここにメモします。

こちらの記事は、Pytorch LightnigでGPUを指定する方法 に関する以下のセッションの内容をまとめたいと思います。

1. pl.Trainerの引数で使用するGPUの定義

pl.Trainerの引数にて、gpus=[使用したいGPUの番号]　とすることで、GPUを指定出来ます。

trainer = pl.Trainer(
                    accelerator="gpu",
                    gpus=[使用したいGPUの番号],
                    )

ここで注意！！
必ずGPUを指定したい場合は リスト形式 [] で括る必要があります。

gpus=[1] とすることで、以下の画像のように一つの指定したGPUを使用することが出来ました。

2. よくある間違い！！

gpus=使用したいGPUの番号、として整数を与えてしまうと問題が生じます。

例としてgpus=2とした場合を考えます。この引数を与えると、GPU=0, 1, 2を使用する設定になります。整数で与えると、その番号以下のGPUを全て使用することになるのでご注意ください。

# GPUの指定に失敗する例
trainer = pl.Trainer(
                    accelerator="gpu",
                    gpus=2,
                    )

まとめ

以上が、Pytorch lightningで指定のGPUを設定する方法です。
これからも気になった点をまとめていこうと思います！