PyTorchでGPU情報を確認（使用可能か、デバイス数など）
PyTorchでGPUの情報を取得する関数はtorch.cuda以下に用意されている。GPUが使用可能かを確認するtorch.cuda.is_available()、使用できるデバイス（GPU）の数を確認するtorch.cuda.device_count()などがある。
- torch.cuda — PyTorch 1.7.1 documentation
- torch.cuda.is_available() — PyTorch 1.7.1 documentation
- torch.cuda.device_count() — PyTorch 1.7.1 documentation
ここでは以下の内容について説明する。
- PyTorchでGPUが使用可能か確認:
torch.cuda.is_available() - PyTorchで使用できるGPU（デバイス）数の確認:
torch.cuda.device_count() - GPUの名称、CUDA Compute Capabilityを取得
- CUDAが使用するGPUを設定: 環境変数
CUDA_VISIBLE_DEVICES
torch.Tensorやモデル（ネットワーク）をCPUからGPUに転送する方法や、実際にGPUが使われているかを簡易的に確認する方法については以下の記事を参照。
本記事のサンプルコードにおけるPyTorchのバージョンは以下の通り。バージョンが異なると仕様が異なる場合もあるので注意。
import torch
print(torch.__version__)
# 1.7.1
PyTorchでGPUが使用可能か確認: torch.cuda.is_available()
PyTorchでGPUが使用可能かどうかはtorch.cuda.is_available()で確認できる。使用できればTrue、できなければFalseを返す。
print(torch.cuda.is_available())
# True
グラフィックボードがあっても、CUDAのドライバなどの準備が整っていないとFalseとなる。
PyTorchで使用できるGPU（デバイス）数の確認: torch.cuda.device_count()
PyTorchで使用できるGPUの数はtorch.cuda.device_count()で取得できる。
print(torch.cuda.device_count())
# 1
デフォルトのGPU番号（インデックス）はtorch.cuda.current_device()で取得できる。0始まり。
print(torch.cuda.current_device())
# 0
GPUの名称、CUDA Compute Capabilityを取得
GPUの名称およびCUDA Compute Capabilityは、それぞれ、torch.cuda.get_device_name()とtorch.cuda.get_device_capability()で取得できる。
print(torch.cuda.get_device_name())
# GeForce GTX 1080 Ti
print(torch.cuda.get_device_capability())
# (6, 1)
torch.cuda.get_device_capability()は(major, minor)のタプルを返す。上の例の場合、Compute Capabilityは6.1となる。
上の例のように引数を省略した場合は、デフォルト（torch.cuda.current_device()が返すインデックス）のGPUの情報を返す。
他のGPUの情報を取得したい場合は、インデックスの数値やtorch.device型のオブジェクト、および、それを表す文字列で指定できる。以下の例はすべて同じGPUを指定している。
print(torch.cuda.get_device_name(0))
# GeForce GTX 1080 Ti
print(torch.cuda.get_device_name(torch.device('cuda:0')))
# GeForce GTX 1080 Ti
print(torch.cuda.get_device_name('cuda:0'))
# GeForce GTX 1080 Ti
torch.deviceについては以下の記事も参照。
CUDAが使用するGPUを設定: 環境変数CUDA_VISIBLE_DEVICES
CUDAが使用するGPUは環境変数CUDA_VISIBLE_DEVICESで設定できる。
例えばCUDA_VISIBLE_DEVICESを0とすると0番のGPUのみ、0, 2とすると0番と2番のGPUのみが使用される。-1とするとすべてのGPUが使用されなくなる。
CUDAに対する設定なので、システムの環境変数を変更した場合は、PyTorchだけでなくTensorFlowなどCUDAを用いる他のアプリケーションにも影響するので注意。
永続的に変更するには、他の環境変数と同じように設定すればよい。例えばUbuntuでは~/.profileなどで設定する。
Pythonのコード中で環境変数を変更するにはos.environを使う。
以下はGPUを使わないように設定した例。GPUが認識されなくなっている。
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
import torch
print(torch.cuda.is_available())
# False
print(torch.cuda.device_count())
# 0
os.environの設定はimport torchの前に行わないと反映されない模様。要注意。