Python: GPU を使う主要なパッケージで CUDA が有効か確かめる方法について

CUDA を利用する Python のパッケージは、多くの場合それ自体のバージョンだけでなく対応している CUDA Toolkit や cuDNN のバージョンまで気にする必要がある。 なんだか環境やコンテナイメージを作るたびに確認する方法や互換性について調べている気がするので、以下のパッケージについてまとめておく。

• PyTorch
• TensorFlow (2.x)
• CuPy
• CuDF

使った環境は次のとおり。

$ lsb_release -a
No LSB modules are available.
Distributor ID: Ubuntu
Description:    Ubuntu 18.04.5 LTS
Release:    18.04
Codename:   bionic
$ uname -r
4.15.0-124-generic
$ python -V       
Python 3.8.6
$ pip list | egrep -ie "^(torch|tensorflow|cupy|cudf) "
cudf                      0.17.0             
cupy                      8.1.0              
tensorflow                2.4.0              
torch                     1.7.1

もくじ

下準備

あらかじめ各種パッケージをインストールして Python の REPL を起動しておく。

$ python

PyTorch

PyTorch では、次のようにする。

>>> import torch
>>> torch.cuda.is_available()
True

CUDA との互換性を調べる方法

最新バージョンについては以下のページを参照する。

pytorch.org

以前のバージョンを調べたいときは以下。

pytorch.org

TensorFlow (2.x)

TensorFlow では以下のようにする。

>>> import tensorflow as tf
>>> len(tf.config.list_physical_devices('GPU')) > 0
True

CUDA との互換性を調べる方法

以下のページを参照する。 日本語のページは翻訳の関係で新しいバージョンが表示されないことがあるため注意する。

www.tensorflow.org

CuPy

CuPy はインポートでエラーにならなければ使える。

>>> import cupy as cp

あえて関数にするとしたら、こんな感じ？

def cupy_is_available():
    try:
        import cupy as cp
        return True
    except ImportError:
        return False

CUDA との互換性を調べる方法

CuPy は CUDA のバージョンごとにバイナリが用意されているので、それを使えば良い。

github.com

CuDF

CuDF も、インポートでエラーにならなければ使えるはず。

>>> import cudf

少なくとも現行の安定バージョン (0.17) では、ソースコードを確認するとインポートするタイミングでチェックが走っているようなので。

github.com

CUDA との互換性を調べる方法

互換性を調べるときは以下のインストールページを確認するのが良いのかな。

rapids.ai

いじょう。

MSI GeForce RTX 3090 VENTUS 3X 24G OC グラフィックスボード VD7357

Linuxで動かしながら学ぶTCP/IPネットワーク入門