CUDAが使えるTensorFlow環境をDocker上に作る

おおとろ

2024年11月28日 10:49

はじめに

備忘録として、Docker上に、NVIDIA GPU（CUDAのほうがメイン）とTensorFlowを使ってモデルをトレーニングする環境を構築する方法について共有します。

対象の読者は、
「ローカルマシンに構築したいけど、
実環境での実装ははばかられると思っている人」
「AI, DL, NNの初学者」です。

Dockerを入れる

WSLとDockerを入れましょう。別に難しくないのでカット。

CUDAとTensorflowのバージョン決定

ここが詰まるポイント1です。
CUDAのバージョンを決定するうえで、そのCUDAのバージョンが、どのTensorflowのバージョンと互換性があるのかどうかを知る必要があります。

さらに問題をややこしくするのが、Dockerで公式に配布されているTensorflowのコンテナの、Tensorflowのバージョンです。適当に調べてTensorflowの最新バージョンと対応しているCUDAを選べばいいのではなく、Dockerコンテナでのバージョンで調べなければなりません。

2024/11/28現在の最新バージョンは、

Tensorflow: 2.13.0
CUDA: 12.2

となっています。
TensorflowはDocker環境を構築するときに呼び出せばいいので、例えば以下のようなDockerfileを記述すればコンテナを作成できます。

# ベースイメージとして TensorFlow GPU イメージを使用（CUDA 12.2 対応）
FROM tensorflow/tensorflow:2.13.0-gpu

# 作業ディレクトリを設定
WORKDIR /workspace

# 必要なパッケージをインストール
RUN apt-get update && apt-get install -y \
    python3-pip \
    python3-dev \
    libgl1-mesa-glx \
    libglib2.0-0 \
    && apt-get clean \
    && rm -rf /var/lib/apt/lists/*

# Pythonライブラリのインストール
RUN pip install --no-cache-dir \
    Pillow \
    numpy \
    tensorflow \
    keras \
    scikit-learn \
    scikit-image

# 環境変数の設定（CUDA関係）
ENV LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH

# コンテナ起動時にシェルを開く設定
CMD ["bash"]

ライブラリなどは必須でないものも含まれていますが、大抵の場合AutoEncoderモデルを使いたいときは画像の学習だと思うのでこういったライブラリはほとんど必須かと思われます。

地味に詰まるポイント、CUDAのバージョン

先ほどのDockerfileでコンテナを作る前にCUDAをインストールしなければならないわけですが、

ここで私がずっと詰まっていた割には全然大したことないを一つ共有します。

それは「インストールされているCUDAのバージョンがわからない」という問題です。

例として12.5をインストールしてしまったとします。
CUDAの12.5に対応したTensorflowのバージョンはないので、ダウングレードしなければならないわけですが、
ここで、CUDAの12.5をアンインストールして、12.2を入れたとします。
そのうえで、ターミナル上で「nvidia-smi」と打っても、
なんとバージョンは「12.5」のままです。しかも再起動しても変わりません。

12.5が消せていないのか、と思いきや、これが実は12.2にちゃんとダウングレードされています。

それを確認する方法が、「nvcc --version」です。
こちらのほうが正確で、優先すべき情報であるらしいです。なぜかは知らないですが。

また、ほかにもインストールされているディレクトリに移動してみる、など様々方法はあります。

ただ一つ覚えておいてほしいのは「nvidia-smiは参考にならない」ということです。「nvcc --version」を使いましょう。

そして、CUDAをインストールしたら、パスも通しましょう。ここら辺はほかのブログでいくらでもあるので書きません。

おまけ：cuDNN

これもまた互換のあるバージョンを探さなければならない面倒なタイプのやつです。
CUDA12.Xに対応したcuDNNは v8.9.7です。（2024/11/28）
これをインストールするには、解凍したファイルを適切な位置に配置する、というインストールです。まあ難しくはないのでカットします。

cuDNNを用いることでどうだ、という実感はあまりないですが、これはどちらかというとDNNのコードを記述する側にメリットがあり、それゆえに実行（学習）速度も向上するというモノなので、利用者にはあまり関係はありません。ただ、少しでも早くなるなら入れたほうがいいです。