はじめに
現在、深層学習ライブラリは以下の3派閥に分かれており、混迷を極めています。
- PyTorch
- Tensorflow
- JAX
特に、自称書籍マニアとして、いろんな技術書に手を出し、遊ばせてもらっている身からすると、PyTorch を使っている本、Tensorflow を使っている本(そしてごくまれにJAX)と多様なライブラリの環境構築に遭遇します。しかし、最近は毎回毎回環境という名の砂場を作るのが面倒になってきました。というわけで、「今日は PyTorch だなぁ～」とか「Tensorflow を使いたい気分」とか「今回こそ JAX と仲良くなるぞ」と気まぐれな自分のために、全載せ環境を構築する流れをまとめたいと思います。ついでに公開することで、同じような境遇でお困りの方、とりあえず万能兵器が欲しいという方の手助けになれば幸いです。
この記事で分かること
- PyTorch・Tensorflow・JAX・Kerasの導入方法
- 各ライブラリから GPU が使えるかどうかの確認方法
- Kerasのバックエンド変更方法
各ライブラリの導入方法
- Python 3.13(miniconda)
仮想環境は miniconda で作ります。特に意味は無いです。
conda create -n envname python=3.13 -y
仮想環境ができたら、以下のコマンドでアクティベート。
conda activate envname
PyTorchのインストール
続いて、PyTorch の導入です。pip でそれっぽい名前を叩いておけば勝手に cuda 付きをインストールしてくれます。
pip install torch torchvision torchaudio
Tensorflowのインストール
Tensorflow は後ろに「gpu使いたいです」と主張しなければなりません。
pip install tensorflow[and-cuda]
JAX（とFlax・Optax）のインストール
JAX は cuda のバージョンを指定する必要があります。
pip install jax[cuda12] flax optax
Kerasのインストール
KerasといえばTensorflowという刷り込みが激しいため、他のライブラリからも利用可能になったという実感が最近までありませんでした。しかし、主要ライブラリに対応してくれているのならば入れ得なので入れておきましょう。
pip install keras keras-cv keras-hub
各ライブラリからGPUが使えるかどうかの確認方法
適当なnotebookを作って以下のコードを実行しましょう。
PyTorch
import torch
print(torch.__version__)
print(torch.cuda.is_available())
Tensorflow
import tensorflow as tf
print(tf.__version__)
print(tf.config.list_physical_devices("GPU"))
JAX
import jax
print(jax.__version__)
print(jax.local_devices())
Kerasのバックエンド変更方法
Kerasはインストール時、デフォルトのバックエンドが Tensorflow になっていますが、これを PyTorch や JAX に変更できます。但し、PyTorch や JAX がインストールされていることが条件です。このとき、バックエンドを変更する方法は主に２つあります。１つはコード上で変更する方法、もう１つはローカルの設定ファイルを書き換える方法です。前者は Google Colab などのクラウド環境で実行するとき、後者は（自分一人が好き勝手出来る）ローカルマシンで実行するときに有効です。
コード上で変更する方法
Pythonのosモジュールを使います。注意したいのは、変更前に Keras を import してはいけない、つまり、バックエンドの変更後に Keras を import しないといけないということです。
import os
os.environ["KERAS_BACKEND"] = "jax"
この後に Keras を import します.
import keras
バックエンドを PyTorch にしたいという場合は、"jax" の部分を "torch" にしてください。
ローカルの設定ファイルを書き換える方法
一度 Keras を import すると、ホームディレクトリに ".keras" が自動生成されます。
.keras ディレクトリには "keras.json" が存在しており、これを書き換えることでバックエンドを変更することができます。
{
...,
# 以下を "tensorflow"・"torch"・"jax" のいずれかにする
"backend": "jax",
...,
}
バックエンドが変更されているかどうかは以下のコードで確認できます。
import keras
print(keras.__version__)
print(keras.backend.backend())
おわりに
「各ライブラリごとに環境作れ」の一言で全て吹き飛ぶような記事ですが、極々稀に PyTorch を使いつつ何故か Tensorflow を要求する意味不明な本に遭遇します。こういう場合には役立つような気がします。
それでは、用途・容量を守って楽しい深層学習ライフを過ごしましょう。
おまけ
どこまで許されるか試してみたくなり、以下の "requirements.txt" を作成し、"pip install -r requirementx.txt" で一括導入しました。結果、以下の場合は各ライブラリで GPU が使えます。すごいですね。
torch
torchvision
pytorch-lightning
pytorch-ignite
timm
tensorflow[and-cuda]
keras
keras-cv
keras-hub
keras-nlp
jax[cuda12]
flax
optax
albumentations
transformers[torch,ja]
diffusers
datasets
bitsandbytes
scikit-learn
matplotlib
japanize-matplotlib
pandas
seaborn
tqdm
opencv-python
jupyter