チューリングのMLOpsエンジニアの岩政です。
先日ポストした内容が思いのほか拡散されたので、個人的な再現性が保ちやすい PyTorchの環境構築の流れ についてまとめた記事です。
本記事の主な対象は、Linux OSでNVIDIA GPUが搭載されているマシンでPyTorchをインストールする方です。他のOSやWSL2などで同様の操作が可能かは保証しません。
本記事で紹介したい流れを具体的なコマンドをもとに先に紹介します。
1.1 NVIDIA Driverのインストール
NVIDIA Driverをインストールしておく必要があります。
Ubuntuであれば、Ubuntu Server documentation – NVIDIA drivers installationなどに従ってインストールします。本記事では、NVIDIA Driverのインストール方法について説明をしません。
1.2 uvのインストール
Pythonのプロジェクト管理ツールであるuvをインストールします。特に理由がなければ、以下のようなコマンドで最新版のuvをインストールします。
curl -LsSf https://astral.sh/uv/install.sh | sh
なお本記事のuvのバージョンは v0.8.13 です。以下のコマンドで確認できます。
1.3 uvでPythonプロジェクトの作成
uv init でプロジェクトディレクトリを作成することができます。torch-project という名前でディレクトリを作成しましょう。
torch-project に移動します。このディレクトリは以下のような構成になっています。
$ tree -a -L 1
torch-project
├── .git
├── .gitignore
├── .python-version
├── README.md
├── main.py
└── pyproject.toml
pyproject.toml でuv用の設定が含まれたら uv sync で仮想環境 .venv を作成することができます。
以下を実行して環境をアクティベートすることができます。今回は実行してもしなくても問題ありません。
1.4 uv pip install を用いたtorchのインストール
以下のコマンドでインストールします。
uv pip install torch --torch-backend=auto
torchの任意のバージョンをインストールする場合は、クォーテーションで=={$version}を指定してください。例えば、以下のコマンドです。
uv pip install "torch==2.7.1" --torch-backend=auto
以下のコマンドでPyTorchがCUDAを利用できるか確認します。True が返ってくることを期待します。
uv run python -c "import torch; print(torch.cuda.is_available())"
1.5 torchのCUDA versionからindexを調べる
torchのCUDAのversionを調べます。以下のコマンドや uv pip list | grep torch で出力されたversionの cu 以降を調べます。以下であれば cu129 です。
Version: 2.8.0+cu129
Location: /workspace/torch-project/.venv/lib/python3.11/site-packages
Requires: filelock, fsspec, jinja2, networkx, nvidia-cublas-cu12, nvidia-cuda-cupti-cu12, nvidia-cuda-nvrtc-cu12, nvidia-cuda-runtime-cu12, nvidia-cudnn-cu12, nvidia-cufft-cu12, nvidia-cufile-cu12, nvidia-curand-cu12, nvidia-cusolver-cu12, nvidia-cusparse-cu12, nvidia-cusparselt-cu12, nvidia-nccl-cu12, nvidia-nvjitlink-cu12, nvidia-nvtx-cu12, sympy, triton, typing-extensions
Required-by:
1.6 pyproject.toml に記載する
uv init で作成された pyproject.toml に以下のように [tool.uv.sources] と [[tool.uv.index]] を記載して uv sync を実行します。 cu128 が出力された方は url を "https://download.pytorch.org/whl/cu128" に設定するなど表示されたCUDA versionに合わせてください。
正しくそのIndex URLが存在するかどうかは PyTorchのGet Started(最新版のtorch)やInstalling previous versions of PyTorchに記載されている --index-url を確認しましょう。
pyproject.toml
[project]
name = "torch-project"
version = "0.1.0"
description = "Add your description here"
readme = "README.md"
requires-python = ">=3.12"
dependencies = [
"torch==2.8.0",
]
[tool.uv.sources]
torch = [
{ index = "torch-cu129" },
]
[[tool.uv.index]]
name = "torch-cu129"
url = "https://download.pytorch.org/whl/cu129"
explicit = true
無事に uv sync を実行できたら完了です。
PyTorchのCUDA実行には、NVIDIA製のCUDA対応GPUとNVIDIA Driverが必要です。正しくドライバーがインストールできれば nvidia-smi でドライバーバージョンやCUDAバージョン、GPUの使用状況が確認できます。
現在、PyTorchをNVIDIA GPUで効率よく利用するにはCUDA Toolkitを追加でインストールする必要はありません。
これについてはKazuki Fujiiさんの[Tips] PyTorchにおける動的リンクを参照してください。torchのインストール時に、CUDA ToolkitやcuDNN, ncclなど関連するライブラリも一緒にインストールされます。
2.2 Pythonのプロジェクト管理ツール: uv
uvは注目を集めているPythonのプロジェクト管理ツールです。依存関係解決が他のパッケージ管理ツールよりも高速で、Pythonのバージョン管理も単一のツールで完結するのが魅力です。
--torch-backend=auto はuv独自の機能です。これについて知る前に、簡単にuvの機能を知りましょう。
2.2.1 pyproject.toml にまとめる
uv init を実行すると様々なファイルが生成されます。その中の1つの pyproject.toml にユーザーはインストールしたいライブラリ名を記述したり、様々なツールの設定を記載することができます。
例えば、Pythonで効率的に数値計算を行うライブラリである numpy をインストールするとします。以下を実行してみましょう。
すると pyproject.toml の dependencies に以下のように追加されます。
dependencies = [
"numpy>=2.3.2",
]
任意のバージョンを指定する場合は、こちらもクォーテーションで指定することができます。また、直接 pyproject.toml を書き換えて、 uv sync することでインストールできます。
このライブラリはどこから取ってきているのでしょうか。基本的には、Python Package Index (PyPI) というPythonライブラリを配布しているリポジトリからインストールすることになります。ここでインストールしたいライブラリを検索することで、公開されているバージョン一覧を確認することができます。
uvの使い方については公式ドキュメントにしっかりとまとめられているので、そちらを参照することをオススメします (https://docs.astral.sh/uv/)。また、v0.3.1とversionが古く一部互換性がないですが、基本的な使い方は以下にまとめています。
2.2.2 uv pip install と uv add どちらを使うべきか
uvのライブラリのインストールは、uv pip install と uv add の2つがあります。どちらを使えばいいでしょうか。
個人的には、基本的には環境構築の再現性が取れやすい uv add を活用した pyproject.toml にまとめていく方法を推奨します。uvは dependencies などに記載したライブラリの依存関係を uv.lock に管理して、 uv sync で高速にインストールします。チーム開発や長期間の個人開発の際に複数の uv pip install で管理するのは手間で、やはり pyproject.toml にまとめた方が意思疎通がしやすく再現性も担保しやすいです。
そのため本記事も、はじめに uv pip install torch --torch-backend=auto を用いて動作可能な torch のCUDAバージョンを調べて、最終的には pyproject.toml にまとめることを推奨します。
2.3 Index URLとは？
ここまで登場していた --index-url とは何でしょうか。
--index-url は、PyPIのようなPythonライブラリを配布しているリポジトリのURLを指定します。pip であればインストール時にはこのURLから指定のライブラリに関するHTMLを取得して、そこからPythonライブラリの配布形式であるwheelなどの最適なリンクを抽出します。詳細は、Finding and choosing filesを参照してください。
uvでも通常はPyPIの https://pypi.org/simple が指定されています。torch も https://pypi.org/project/torch/ に公開されていますが、これは1つのライブラリバージョンに対してCUDAバージョンも1つのもののみ公開されています(2025年8月29日現在、torch 2.8.0 はおそらく cu128 です)。
この --index-url を変えるために uv pip であれば --extra-index-url を設定することができます。例えば cu129 が指定された方は、以下のコマンドで同じ torch がインストールされます。
uv pip install torch --extra-index-url https://download.pytorch.org/whl/cu129
2.3.1 --torch-backend=auto は何をする？
uv pip install torch --torch-backend=auto は、この torch の --extra-index-url を適切なIndex URLに自動で選択するものです。有効にすることで、CUDA Driverなどの有無を確認して、関連するライブラリに対して最も互換性の高いPyTorch Indexを使用します。
これはAutomatic backend selectionで簡単に説明されています。また、これ自体は v0.6.9 からのpreview機能で、uv pip install のみでサポートしています。
この --extra-index-url を pyproject.toml を明示的に記載することができます。
方法として、まず [[tool.uv.index]] を記載します。name は任意の文字列で、url にPyTorch用のURLを、 explicit = true を設定します。
[[tool.uv.index]]
name = "torch-cu129"
url = "https://download.pytorch.org/whl/cu129"
explicit = true
次に [tool.uv.sources] を設定します。ライブラリごとにどのindexからインストールするかを指定します。なお explicit = true を設定することにより、[tool.uv.sources] で指定しない場合はそのURLを使わずにデフォルトのPyPIのURLを使用します。
[tool.uv.sources]
torch = [
{ index = "torch-cu129" },
]
ここでは紹介しませんが [tool.uv.sources] には marker 機能が存在します。OSやCPUアーキテクチャごとにインストールするindexを変更することが可能で、特にmacOSでCPU版の torch を指定する際に便利です(PEP 508)。以下の記事でその詳細を説明しています。
このように記載した pyproject.toml が存在するプロジェクトで uv sync を実行することで、適切なCUDAバージョンの torch をインストール可能です。なお、PyTorchのインストール方法のベストプラクティスについて Using uv with PyTorch で紹介されています。
この記事では uv pip install torch --torch-backend=auto を用いた互換性の高い torch のインストールから、再現性を担保するための pyproject.toml に相当するIndexを指定するための方法を紹介しました。
なお、そもそも uv pip install torch --torch-backend=auto で適切な torch がインストールできない可能性に十分注意すべきです。uvのissueにも、このコマンドがうまくいかない場合も言及されておりその解決は個々人の環境で頑張る必要があります。応援しています。
Views: 0
