uv を活用して CUDA 対応の Python 環境を構築するとき、つまづく部分があるので方法を記載します

CUDA環境のインストール

  1. 最新のドライバをインストール

  2. CUDA Toolkit をインストール

確認用コマンド:

1
2
3
4
5
6
7
8
# GPUの状態を確認
nvidia-smi

# インストール済みの CUDA バージョンを確認
nvcc --version

# CUDA_PATH のパスを確認
Get-ChildItem Env:CUDA_PATH

プロジェクトの作成

  1. プロジェクトを新規作成
1
2
3
uv init my-torch-project
cd my-torch-project
uv python pin 3.12
  1. PyTorch の CUDA 対応バージョンをインストール

このまま uv add torch==2.5.1+cu124 を実行すると、No solution found when resolving dependenciesとなります

これを回避するため、以下の設定を pyproject.toml に追加します

1
2
3
4
5
6
7
8
9
10
11
12
13
14
[[tool.uv.index]]
name = "torch-cuda"
url = "https://download.pytorch.org/whl/cu124"
explicit = true

[[tool.uv.index]]
name = "torch-cpu"
url = "https://download.pytorch.org/whl/cpu"
explicit = true

[tool.uv.sources]
torch = [
  { index = "torch-cuda" },
]
  1. 設定完了後に、以下のコマンドを実行
1
uv add torch==2.5.1+cu124

CUDAの動作確認

テスト用コード hello.py:

1
2
3
4
5
6
7
8
9
10
11
import torch

def main():
    cuda_available = torch.cuda.is_available()
    print(f"CUDA Available: {cuda_available}")
    if cuda_available:
        cuda_version = torch.version.cuda
        print(f"CUDA Version: {cuda_version}")

if __name__ == "__main__":
    main()

上記のコードを実行

1
uv run hello.py

実行結果に、CUDAの利用可否とバージョンが表示されます