Zenn
📘

Dockerを用いてGPU版TensorFlow と Jupyter Notebook を動かしてみた

2024/04/19に公開
Docker
Jupyter Notebook
TensorFlow
tech

概要

TensorFlowなどの大きなパッケージをそのままインストールすると管理が面倒なことが多いので、これをDockerで管理して実行することを目指しました。

前提条件

https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html

上記のURLの内容に従ってNVIDIA Toolkitを入れておく必要があります。

GPU版TensorFlowを使う

Dokcerfileの先頭でGPU版のTensorFlowを指定してあげればよいだけです。

Dockerfile
FROM tensorflow/tensorflow:latest-gpu

次に、今回はあまり説明しませんが、今後必要になりそうな(遊べそうな) TransformersとTorchをインストールするように指示をします。

Dockerfile
RUN pip3 install --upgrade pip
RUN pip3 install Torch transformers

Jupyter Notebook を導入する

Jupyter Notebookを使用すれば色々簡単にできると思い同居させることとにしました。

Dockerfile
RUN pip3 install jupyter notebook

Jupyter NotebookはNon-Dockerの構成を想定しているので、このままだと動きません。
なぜならDokcer内のlocalhostでListenしてしまうので、これを0.0.0.0でListenするようにしなければならないのです。

そこでconfigファイルを一筆書きます。

./py/jupyter_notebook_config.py
c.NotebookApp.ip = '0.0.0.0'

そのconfigファイルの読み込みと、rootでの実行を許可するオプションをつけて実行すれば良のです。

Dockerfile
CMD ["jupyter-notebook","--allow-root","--config=/py/jupyter_notebook_config.py"]

まとめると以下になります。

Dockerfile
FROM tensorflow/tensorflow:latest-gpu
RUN pip3 install --upgrade pip
RUN pip3 install Torch transformers
RUN pip3 install jupyter notebook
CMD ["jupyter-notebook","--allow-root","--config=/py/jupyter_notebook_config.py"]

実行

あとは、これをDocker Composeから呼べるようにすれば良いわけです。

compose.yml
version: "3"
services:
  tensorflow:
    build: .
    runtime: nvidia
    environment:
      - NVIDIA_VISIBLE_DEVICES=all
    volumes:
      - ./py:/py
    ports:
      - 8888:8888

そのあとlogに出てくる(ここでは最下の)URLをたたけばOKです。

docker compose logs
tensorflow-1  |
tensorflow-1  |     To access the server, open this file in a browser:
tensorflow-1  |         file:///root/.local/share/jupyter/runtime/jpserver-1-open.html
tensorflow-1  |     Or copy and paste one of these URLs:
tensorflow-1  |         http://967b100b446b:8888/tree?token=cb549417905d0e8d9f145ed13a6ab3f6ac3c10fe57fe3180
tensorflow-1  |         http://127.0.0.1:8888/tree?token=cb549417905d0e8d9f145ed13a6ab3f6ac3c10fe57fe3180

念のためGPUが使えるか確認します。

$ docker compose exec tensorflow bash

________                               _______________
___  __/__________________________________  ____/__  /________      __
__  /  _  _ \_  __ \_  ___/  __ \_  ___/_  /_   __  /_  __ \_ | /| / /
_  /   /  __/  / / /(__  )/ /_/ /  /   _  __/   _  / / /_/ /_ |/ |/ /
/_/    \___//_/ /_//____/ \____//_/    /_/      /_/  \____/____/|__/


WARNING: You are running this container as root, which can cause new files in
mounted volumes to be created as the root user on your host machine.

To avoid this, run the container by specifying your user's userid:

$ docker run -u $(id -u):$(id -g) args...
# python3.11
Python 3.11.0rc1 (main, Aug 12 2022, 10:02:14) [GCC 11.2.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> from tensorflow.python.client import device_lib
>>> print(device_lib.list_local_devices())
[name: "/device:CPU:0"
device_type: "CPU"
memory_limit: 268435456
locality {
}
incarnation: 15791178877055289296
xla_global_id: -1
, name: "/device:GPU:0"
device_type: "GPU"
memory_limit: 1249247232
locality {
  bus_id: 1
  links {
  }
}
incarnation: 942658034812053152
physical_device_desc: "device: 0, name: NVIDIA GeForce RTX 3060, pci bus id: 0000:01:00.0, compute capability: 8.6"
xla_global_id: 416903419
]

まとめ

IPアドレスの設定を編み出すのに時間がかかりました。これができないときちんとJupyter Notebookが使えないですが、この方法をマスターしておくと後で応用ができ、楽になると考えられます。

Discussion