はじめに

現在、深層学習ライブラリは以下の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" が存在しており、これを書き換えることでバックエンドを変更することができます。

~/.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 が使えます。すごいですね。

requirements.txt

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