Tensorflow2.x + JupyterLab Docker環境を作る

Abstract

通常、Tensorflowの開発環境はJupyter Notebookですが、最近はJupyter Labがデファクトスタンダードになりつつある感じが個人的にあります。
自分のリマインダーも兼ねて、Docker設定周りをまとめてみます。
（Docker/docker-composeの導入はご自分の環境に合わせて予め導入をお願いします）

Results

希望するディレクトリに以下のDockerfileとdocker-compose.ymlを作成します。

Dockerfile

Dockerfile

# python3をベースにする
FROM tensorflow/tensorflow:latest-devel-py3

RUN apt-get update && apt-get install -y \
    curl \
    git \
    tree \
    vim
# RUN pip install --upgrade pip
RUN pip3 install --upgrade pip setuptools
RUN pip3 install jupyterlab
RUN pip3 install matplotlib # デフォルトのTensorflowにはmatplotlibがない！！ <- 実際に遊んでいて初めて気付きました。。
RUN jupyter serverextension enable --py jupyterlab

# 作業ディレクトリを指定
WORKDIR /workdir

# コンテナ側のポート番号
EXPOSE 8888

ENTRYPOINT ["jupyter-lab", "--ip=0.0.0.0", "--port=8888", "--no-browser", "--allow-root", "--NotebookApp.token=''"]

# Jupyterlabのルートとなるディレクトリを指定
CMD ["--notebook-dir=/workdir"]

docker-compose.yml

docker-compose.yml

version: '3'
services:
  dev:
    build:
      context: .
      dockerfile: Dockerfile
      # container_name: tf-container
    ports:
    - "8888:8888"
    volumes:
    - ./workdir:/workdir

起動

同ディレクトリ階層で下記のコマンドを実行します。

$ docker-compose up -d

初回は基盤イメージのDLから行われるので時間がかかります（作業ログがterminal内で表示される）が、次回以降はすぐに起動します。
無事に起動したら、localhost:8888でJupyter Labにアクセスできます。

上記の方法では、/workdirがローカルストレージと接続しているため、container / imageを後で削除しても作成したファイルは残ります。
後日、gpu対応に変更したり、TensorboardをJupyter Lab Extensionsで結合する処理をDockerfile内に追記してコンテナを再作成する等、行う場合も気兼ねなく環境を再編することができるので便利だと思います。

他にも良い書き方があると思います。Better ideaあれば教えて頂けると嬉しいです。