【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

https://imagingsolution.net/deep-learning/pytorch/pytorch_mnist_sample_program/

- TensorFlow

https://qiita.com/exy81/items/bcfe680f5e04098dbc6b

# 環境作成方法

- Dockerfile

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

```dockerfile: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
```

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

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

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

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

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

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

```shell:TensorFlow
pip install tensorflow[and-cuda]
```

# 参照したもの

- PyTorch

https://pytorch.org/

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3568254/e0cd1d23-2a4f-6594-8ea8-321d6b36bfdd.png)

- TensorFlow

https://www.tensorflow.org/install/source#gpu_support_2

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3568254/0541c24c-abe2-6fd2-3bf7-10f339526c63.png)

https://www.tensorflow.org/install/pip