WSLとDockerを使ってTensorFlow GPU環境を構築
はじめに
背景
もともとWSLにpipでTensorFlowをインストールして使っていました。最近はpip install tensorflow[and-cuda]のコマンド一発でGPUが使えるので便利ですね。（執筆時点）
しかし、TensorFlow以外でcuda依存のソフトを使いたくて、cudaのインストールとかごにょごにょしていたら、いつの間にかTensorFlowがGPUを認識しなくなっていました😭。（何をやらかしたのか記憶にない...）
WSLやグラフィックドライバの再インストールはしてみましたが、それでも認識しません。諦めて、Docker上でTensorFlowを動かすことにしました。この記事は、その手順の備忘録です。
概要
この記事では、WSLとDockerを使ってTensorFlowをGPU上で動かす方法を紹介します。また、VSCodeからDockerコンテナにアクセスする方法も記載します。
目標
- TensorFlowをGPU上で動かす！（最重要）
- WSLのディレクトリをDockerコンテナにマウントして、Jupyter Notebookを編集・実行したい。
- VSCodeからWSLに接続する感じで、Dockerコンテナに接続したい。
- VSCodeがPython環境を正しく認識して自動補完などが効くようにしたい。
実行環境
- OS: Windows11 23H2
- GPU: GeForce RTX4060Ti 16GB
- WSLのバージョン:
> wsl --version
WSL バージョン: 2.1.5.0
カーネル バージョン: 5.15.146.1-2
WSLg バージョン: 1.0.60
MSRDC バージョン: 1.2.5105
Direct3D バージョン: 1.611.1-81528511
DXCore バージョン: 10.0.25131.1002-220531-1700.rs-onecore-base2-hyp
Windows バージョン: 10.0.22631.3296
Dockerのインストール
WindowsにDockerをインストールします。
インストールが完了したら、WindowsでDocker Desktopを起動します。
これで、WSL上でdockerコマンドが使えるようになっているはずです。
$ docker -v
Docker version 25.0.3, build 4debf41
$ docker
The command 'docker' could not be found in this WSL 2 distro.
We recommend to activate the WSL integration in Docker Desktop settings.
For details about using Docker Desktop with WSL 2, visit:
https://docs.docker.com/go/wsl2/
DockerでTensorFlow環境を構築
こちらにインストール方法が記載されています。
以下は私好みのオプションで実行しています。
まずはdocker pullでイメージをダウンロードします。
$ docker pull tensorflow/tensorflow:latest-gpu-jupyter
続いてコンテナを作成します。WSL上にあるノートブックを実行したいので、-vオプションでディレクトリをマウントします。--nameでつけるコンテナ名は今回tf_gpuとしていますが、好きな名前にしてください。
$ docker run --gpus all -it --name tf_gpu -v [WSL上のパス]:[Docker上のパス] tensorflow/tensorflow:latest-gpu
Dockerのbashが起動します。以下のような画面が表示されると思います。
これでDockerのコンテナが起動した状態になります。続いて、VSCodeからコンテナにアクセスするための設定を行います。
VSCodeでDockerコンテナにアクセス
VSCodeはこちらからダウンロードできます。Cursorでもいいです。
VSCodeを起動したら、Dev Containersという拡張機能をインストールします。これを入れるとDockerコンテナにアクセスできます。
拡張機能がインストールできたら、左下の><をクリックして、Attach to Running Containerを選択します。
すると、起動中のコンテナ（例の場合だとtf_gpu）が表示されるので、それを選択します。
好きなノートブックを開いたら、右上のSelect Kernelをクリックして、適当なPython環境を選択します。
よくわからないのでRecommendedを選びました。
ノートブックのセルを実行しようとすると.ipykernelパッケージが必要だと言われるので、installをクリックします。
無事、ノートブックが実行できるようになりました。GPUも認識されているようです。
Dockerコンテナの起動・停止
先ほどはdocker runでコンテナを起動して、シェルに入った状態で放置していました。シェルから抜けるにはexitコマンドを使います。exitするとコンテナは停止します。
コンテナの状態確認
docker ps -aで、コンテナの状態を確認できます。
$ docker ps -a
CONTAINER ID IMAGE COMMAND CREATED STATUS PORTS NAMES
1b552472c3fe tensorflow/tensorflow:latest-gpu "/bin/bash" 6 hours ago Exited (137) 2 seconds ago tf_gpu
コンテナの起動
$ docker start tf_gpu
コンテナの停止
$ docker stop tf_gpu
コンテナの削除
$ docker rm tf_gpu
コンテナのシェルに入る
$ docker exec -it tf_gpu bash
Docker Desktopからコンテナを操作する
WSLでDockerのコンテナを作成したら、Docker Desktopからも操作できます。
Jupyter Notebookサーバを立てる
VSCodeだけでなく、ブラウザからもノートブックを開きたいときがあるかもしれません。その場合は、Jupyterも含んだイメージが便利です。
$ docker run --gpus all -it -p 8888:8888 --name tf_gpu_jupyter -v [WSL上のパス]:[Docker上のパス] tensorflow/tensorflow:latest-gpu-jupyter bash
普通に走らせると勝手にJupyterが立ち上がります。個人的にはそれがいやなので、bashでシェルに入ってから、以下のコマンドで好きなときにJupyterを立ち上げるようにしています。
$ jupyter notebook --ip 0.0.0.0 --no-browser --allow-root
http://127.0.0.1:8888/?token=[トークン]というURLが表示されるので、ブラウザでアクセスします。
こちらもGPUが認識されました！
まとめ
どうにかTensorFlowをGPUで動かすことができました。Dockerを使ったのは今回がほぼ初めてなので、無駄な手順とかあるかもしれません。もっとスマートな方法があれば教えていただけると嬉しいです。
Discussion