【CUDA】PyTorchとTensorFlowが共存できる環境構築(2024.12)

PyTorchとTensorFlowで要件とされるCUDAバージョンが異なっていたりするので、CUDA、PyTorch、TensorFlowそれぞれのバージョンの組み合わせで同一のDockerコンテナで動作可能なものを試してみました。
(今回は簡単に1つのコンテナにつき1つのCUDAバージョンで行きます)

共通環境

• WSL2
• Dockerコンテナ上での動作
• Ubuntu24.04
• NvidiaDriver565.90

結論(動作確認できたもの)

CUDAバージョン	PyTorch	TensorFlow	備考
12.5	PyTorch2.5.1(cu124), Python3.12.3	TensorFlow2.18.0, Python3.12.3
12.4	PyTorch2.5.1(cu121), Python3.12.3	TensorFlow2.18.0, Python3.12.3	venv等で環境分離必要
12.3	PyTorch2.5.1(cu121), Python3.12.3	TensorFlow2.18.0, Python3.12.3	venv等で環境分離必要
12.2	PyTorch2.5.1(cu121), Python3.12.3	TensorFlow2.18.0, Python3.12.3	venv等で環境分離必要
12.1	PyTorch2.5.1(cu121), Python3.12.3	TensorFlow2.18.0, Python3.12.3	venv等で環境分離必要
11.8	PyTorch2.5.1(cu118), Python3.12.3	TensorFlow2.18.0, Python3.12.3

(特にTensorFlow公式サイト記載のバージョン対応表で対応していない組み合わせでも動作しているような挙動が見られたので、結果はご参考まで)

なお、動作確認は以下サイトのコードを使用した。

• PyTorch

• TensorFlow

環境作成方法

• Dockerfile

各CUDAバージョンを持つ環境は以下のDockerfileによって作成。
作成の際はDockerfile中の12.5のところを任意のバージョンに変更する。

dockerfile

FROM ubuntu:24.04

# aptの更新など
RUN apt-get update
RUN apt-get upgrade -y
RUN apt-get install -y wget

# CUDAのインストール
RUN wget https://developer.download.nvidia.com/compute/cuda/repos/wsl-ubuntu/x86_64/cuda-keyring_1.1-1_all.deb
RUN dpkg -i cuda-keyring_1.1-1_all.deb
# ↓Ubuntu23.04以降でCUDAのインストールに必要なコマンド
RUN echo "deb http://archive.ubuntu.com/ubuntu/ jammy universe" >> /etc/apt/sources.list.d/jammy.list
RUN echo "Package: *\nPin: release n=jammy\nPin-Priority: -10\n\nPackage: libtinfo5\nPin: release n=jammy\nPin-Priority: 990" >> /etc/apt/preferences.d/pin-jammy
# ↑Ubuntu23.04以降でCUDAのインストールに必要なコマンド
RUN apt-get update
RUN apt-get -y install cuda-toolkit-12-5
RUN rm cuda-keyring_1.1-1_all.deb
RUN echo "export PATH=/usr/local/cuda-12.5/bin\${PATH:+:\${PATH}}" >> /root/.bashrc
RUN echo "export LD_LIBRARY_PATH=/usr/local/cuda-12.5/lib64:\${LD_LIBRARY_PATH:+:\${LD_LIBRARY_PATH}}\n"  >> /root/.bashrc

# .bashrcの編集
RUN sed -i 1iforce_color_prompt=yes /root/.bashrc
RUN echo alias la="'ls -lah'\n" >> /root/.bashrc

# pyenvのインストール
RUN apt-get install -y curl git build-essential libssl-dev libffi-dev libncurses5-dev zlib1g zlib1g-dev libreadline-dev libbz2-dev libsqlite3-dev make gcc liblzma-dev python3-tk tk-dev
RUN curl https://pyenv.run | bash
RUN echo "export PYENV_ROOT=\"\$HOME/.pyenv\""  >> /root/.bashrc
RUN echo "export PATH=\"\$PYENV_ROOT/bin:\$PATH\""  >> /root/.bashrc
RUN echo "eval \"\$(pyenv init -)\""  >> /root/.bashrc
RUN echo "eval \"\$(pyenv virtualenv-init -)\""  >> /root/.bashrc

• イメージの作成コマンド

docker build -f dockerfile -t cuda_ubuntu:12.5-24.04 .

• コンテナの作成と起動のコマンド

docker run -itd --gpus all --name cuda125_ubuntu24 -w /work cuda_ubuntu:12.5-24.04

• コンテナ内でvenv環境を作成し、以下コマンドでPyTorchやTensorFlowをインストール

PyTorch

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cuXXX
# XXXの部分は指定のCUDAバージョンに対応させる

TensorFlow

pip install tensorflow[and-cuda]

参照したもの

• PyTorch

• TensorFlow