🔖

任意のCUDAバージョンのPyTorchをPoetryでインストールする

2022/10/26に公開

Python

PyTorch

poetry

tech

概要

PyTorchからGPUを利用するためには、CUDAのバージョンに合ったパッケージをインストールする必要があります。PyTorchの"Get Started"ページでは、プラットフォームやバージョン、パッケージマネージャーを選択すると、それに応じたインストールコマンドが表示される仕組みがありますが、ここにはPoetryは含まれていません。

そこでこの記事では、Poetryを利用しているプロジェクトにおいて、任意のCUDAバージョンのPyTorchパッケージをインストールする方法を解説します。

方法

実現する方法はいくつか存在するので、それぞれ解説していきます。ここでは、例としてCUDA 11.6対応版である1.12.1+cu116を対象に説明します。

ちなみに、私がおすすめする方法は 1 > 2 > 3 の順です （理由は後述）。

実行環境
- Poetry (version 1.2.2)
- Ubuntu 22.04.1 LTS

1. pyproject.tomlに直書きする

Poetryがパッケージのバージョンや依存関係を管理するpyproject.tomlに、インストールしたいPyTorchの情報を直書きします。

[tool.poetry.dependencies]に、インストールしたいwhlファイルのURLを以下のように指定します。環境に合わせたwhlファイルは、pipの--extra-index-urlで指定するURL先から探すことができます。cp39はPythonのバージョン、cu116はCUDAのバージョンを表しています。

[tool.poetry.dependencies]
torch = {url = "https://download.pytorch.org/whl/cu116/torch-1.12.1%2Bcu116-cp39-cp39-linux_x86_64.whl"}

その後、通常通りpoetry installを行うと、指定したパッケージがインストールされます。

$ poetry install

実際にimport torchでバージョンを確認すると、きちんとCUDA 11.6対応版がインストールされていることがわかります。

In [1]: import torch

In [2]: torch.__version__
Out[2]: '1.12.1+cu116'

2. Package Sourcesを追加する

Poetryにはインストール元にPyPIを使う他に、独自のレポジトリを追加する仕組みがあります。これにPyTorchの任意のバージョンのレポジトリのURLを指定することで、そこからインストールすることが可能です。なお、この方法はPoetry 1.2.0以上である必要があります。

まずsourceにレポジトリのURLを追加します。sourceには名前を付与できるので、わかりやすい名前を付けておきます。

$ poetry source add torch_cu116 https://download.pytorch.org/whl/cu116/
Adding source with name torch_cu116.

そして、通常通りpoetry add で何も指定せずtorchをインストールします。

$ poetry add torch
Creating virtualenv poetry-pytorch-test in /home/yag_ays/poetry_pytorch_test/.venv
Using version ^1.12.1+cu116 for torch

Updating dependencies
Resolving dependencies... (0.8s)<c1>Source (torch_cu116):</c1> Authorization error accessing https://download.pytorch.org/whl/cu116/typing-extensions/
Resolving dependencies... (1.0s)<c1>Source (torch_cu116):</c1> Authorization error accessing https://download.pytorch.org/whl/cu116/typing-extensions/
Resolving dependencies... (1.1s)

Writing lock file

Package operations: 2 installs, 0 updates, 0 removals

  • Installing typing-extensions (4.4.0)
  • Installing torch (1.12.1+cu116)

これでインストールは完了です。先ほどと同様に、torch.__version__で1.12.1+cu116が入ることを確認できます。

ちなみに、pyproject.tomlには以下のように記述されます。

[[tool.poetry.source]]
name = "torch_cu116"
url = "https://download.pytorch.org/whl/cu116/"
default = false
secondary = false

このようにPoetryの機能を利用し、明示的にURLを指定しなくてもインストールが可能なPackage Sourcesですが、一方で以下のような問題も同時に生じます。

注意1： Poetryの実行にかなり時間が掛かる

Package Sourcesを利用すると、明示的にPyTorchのパッケージが依存関係の管理に含まれます。その際には利用環境に対応したパッケージ以外にも、関連するすべてのパッケージが一度ダウンロードされます。そのため、torchのようなパッケージそのもののファイルサイズが大きいものを対象にした場合には、Poetryの実行にかなりの時間がかかるため注意が必要です。

Why does Poetry insist on downloading all candidate distributions for all platforms when metadata is not available?
The need for this stems from the fact that Poetry’s lock file is platform-agnostic. This means, in order to resolve dependencies for a project, Poetry needs metadata for all platform specific distributions. And when this metadata is not readily available, downloading the distribution and inspecting it locally is the only remaining option.
https://python-poetry.org/docs/repositories/#simple-api-repository

注意2： PyPIからパッケージをインストールする際にも登録したPackage Sourcesを見に行く

Package Sourcesを追加した状態で他のパッケージをインストールすると、本来ならPyPIにしかないパッケージに関しても追加したPackage Sources先に存在するかを見に行きます。そのため、下記のようなエラーが大量に出ることになります。ひとしきりエラーが表示された後に通常通りPyPIからインストールされるため、結果としては問題ないのですが、パッケージを追加するたびこれを目にすることになるためかなり煩わしいと感じます。

$ poetry add ipython
Source (torch_cu116): Authorization error accessing https://download.pytorch.org/whl/cu116/ipython/
Using version ^8.5.0 for ipython

Updating dependencies
Resolving dependencies... (0.6s)Source (torch_cu116): Authorization error accessing https://download.pytorch.org/whl/cu116/ipython/
Resolving dependencies... (0.7s)Source (torch_cu116): Authorization error accessing https://download.pytorch.org/whl/cu116/ipython/
Resolving dependencies... (1.4s)Source (torch_cu116): Authorization error accessing https://download.pytorch.org/whl/cu116/colorama/
Resolving dependencies... (2.0s)Source (torch_cu116): Authorization error accessing https://download.pytorch.org/whl/cu116/appnope/
Resolving dependencies... (2.6s)Source (torch_cu116): Authorization error accessing https://download.pytorch.org/whl/cu116/pexpect/
Resolving dependencies... (3.3s)Source (torch_cu116): Authorization error accessing https://download.pytorch.org/whl/cu116/traitlets/
Resolving dependencies... (3.9s)Source (torch_cu116): Authorization error accessing https://download.pytorch.org/whl/cu116/stack-data/
Resolving dependencies... (4.5s)Source (torch_cu116): Authorization error accessing https://download.pytorch.org/whl/cu116/pygments/
Resolving dependencies... (4.7s)Source (torch_cu116): Authorization error accessing https://download.pytorch.org/whl/cu116/prompt-toolkit/
Resolving dependencies... (5.4s)Source (torch_cu116): Authorization error accessing https://download.pytorch.org/whl/cu116/pickleshare/
Resolving dependencies... (6.0s)Source (torch_cu116): Authorization error accessing https://download.pytorch.org/whl/cu116/matplotlib-inline/
[...これが依存するパッケージ全部に及ぶ...]

私が試した限りでは、PyPIをPackage Sourcesに追加したり、PyTorchの方をsecondary = falseに指定したり、poetry add --sourceで指定しても、結局上記のエラーを回避することは出来ませんでした。PyTorch以外のパッケージにもこのような影響が出てしまうのであれば、Package Sourcesを追加せずURL直指定でインストールしたほうが良いかもしれません。

3. Poetryからpipを実行する (非推奨)

他には、Poetryからpipを実行してインストールする方法もあります。これなら、PyTorchのページから確認できるpipコマンドをそのまま利用できます。

$ poetry run pip install torch --extra-index-url https://download.pytorch.org/whl/cu116

ただしこの方法ではpyproject.tomlで管理がされないため、依存関係の管理やコード上での確認が不可能なため、あまり利用しない方がよいでしょう。

まとめ

この記事では、任意のCUDAバージョンに対応したPyTorchパッケージを、Poetryでインストールする方法を解説しました。Poetry 1.12．0よりPackage Sourcesを使うことでpoetry add torchというシンプルな方法でインストールすることができるようになりましたが、依然として依存関係まわりで不便な部分が多く、結局URLやらwhlやらを探すなど色々と面倒な状況は続くようです。