Pythonで生成AIのモデル利用時にGPUを指定する

🔳事前準備

まずは、CUDA Toolkitをインストールします。

CUDA Toolkit
https://developer.nvidia.com/cuda-toolkit

🔳確認

PyTorchのバージョンとの兼ね合いにより、CUDA Toolkitをインストールしただけでは動作しないことがあります。以下の関数を実行して、GPUが利用可能かどうかを確認してください。

torch.cuda.is_available()

利用可能な場合はtrueが返却されます。falseになる場合は、PyTorchのバージョンを確認し、必要に応じて合わせてみてください。

sample code

# GPUが利用可能か確認
if torch.cuda.is_available():
    device = torch.device("cuda")
    print("GPUを使用します。")
else:
    device = torch.device("cpu")
    print("GPUは利用できません。CPUを使用します。")

🔳PyTorchのバージョンを合わせる

GPUを利用するためには、CUDAをサポートするPyTorchをインストールする必要があります。PyTorchのインストール時にCUDAバージョンを指定しているか確認してください。例えば、

pip install torch==2.2.1 torchvision==0.17.1 torchaudio==2.2.1 --index-url https://download.pytorch.org/whl/cu121

のようにインストールする際に、適切なCUDAバージョンを指定します。