スポンサーリンク

PyTorchでGPUを使用するためのtorch.cuda.is_available

Machine Learning

PyTorchでGPUを使用するための確認方法とモデルやTensorへの適用方法について見てみました。

torch.cuda.is_available – GPU使用可能か確認

PyTorchのtorch.cudaでは、CUDAやGPUに対する関数が用意されており、torch.cuda.is_availableでGPUが使用可能か見ることができる。

import torch

print(torch.cuda.is_available())
# True

GPUが使用可能であればTrue、不可能であればFalseが返ってくる。

PyTorchでは、モデル学習時にモデルとデータを同じデバイス(CPU/GPU)に置く必要がある。GPU/CPUの使用によってコードが使えなくなるのを防ぐため、以下コードのように使用可能なデバイスを変数に入れておくと便利です。

device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using {device} device")
# Using cuda device

x = x.to(device) # データをGPU/CPU

net.to(device)   # モデルをGPU/CPU

参考記事、参考資料

PyTorchの入門書で、GPUの利用方法、ネットワークの構築方法や転移学習まで幅広く書かれていてオススメです。

コメント