nvidiaのGPUをDockerのTensorflowで動かす

機械学習を行うためのマシンでnvidiaのGPU RTX-3060+Intel Core i5のマシンを購入しました。そこにLinuxをインストルしてTensorflowで機械学習をする環境を整えたときに大変だったので、手順をブログにまとめます。

作成した開発環境はこのとおりです。

専用マシンなのでDockerなんていらないと思っていたのですが、DockerなしだとGPUがまともに動かず、仕方なくDockerを導入しました。

1. Linux カーネルのダウンロードとインストール

こちらのサイトからUbuntu Desktop 22.04LTSをダウンロードして導入しました。

ダウンロードしたISOファイルをUSBメモリ等に書き込みをします。

ISOファイルの書き込みは、Balena EtcherWin32DiskImagerを利用して書き込みを行うことで、USBメモリ等をブートイメージにすることができます。

作成したUSBをPCに挿入してLinuxをインストールするPCに接続して起動してOSのインストールをしてください。

2. UbuntuのGUIが使いづらい

UbuntuのGUIでも気にならない人は飛ばしてください。

UbuntuのGUIは個人的に好きではありません。私は以前 Lubuntuを使用していたのですがどうも安定しませんでした。今回は結果としてLXDEを導入することで安定動作することが確認できました。

LXDEのGUI
UbuntuのGUI

3. nvidiaのドライバをインストールする

3-1. ドライバのインストール

こちらのリンクからnvidiaのドライバをダウンロードしてインストールをします。

最新の NVIDIA 公式ドライバーをダウンロード

ダウンロードしたファイルを実行してドライバをインストールします。

sh NVIDIA-Linux-x86_64-515.48.07.run

3-2. nvidia dockerのインストール

こちらのサイトを参照してnvidia-docker2をインストールします

4. GPUに対応したTensorflowのDockerイメージファイルをインストール

利用するtensorflowのdockerイメージをダウンロード

docker pull tensorflow/tensorflow:latest-gpu-jupyter

以下のコマンドを入力して動作確認をします。

docker run --gpus all -it tensorflow/tensorflow:latest-gpu-jupyter bash

--gpus allはGPUを全て使用するオプションになります。

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


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...

続いて、nvidia-smiのコマンドを実行して、以下のような画面が表示されたら正しくnvidiaのドライバがインストールされています。

root@5373f491de28:/# nvidia-smi
Mon Jun 27 11:46:24 2022       
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 510.73.05    Driver Version: 510.73.05    CUDA Version: 11.6     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  NVIDIA GeForce ...  Off  | 00000000:01:00.0  On |                  N/A |
|  0%   56C    P8    10W / 170W |     52MiB / 12288MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
                                                                               
+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
+-----------------------------------------------------------------------------+

dockerを抜けるときは、exitを入力してください。そして、docker ps -aqで起動中のdockerのプロセスを表示することができ、docker rmでプロセスをKILLすることができます。

root@5373f491de28:/# exit
docker ps -aq
ca3814cc4195
docker rm ca3814cc4195

5. VS Codeのインストール

開発環境はVisual Studio Code(以下、VS Code)がおすすめです。まず、こちらのサイトにアクセスしてLinux側にVS Codeをインストールします。

導入したOSはUbuntuなので、.debパッケージのダウンロードをします。

Download Visual Studio Code - Mac, Linux, Windows

Visual Studio Code is free and available on your favorite platform - Linux, macOS, and Windows. Download Visual Studio Code to experience a redefined code edi…

ダウンロードしたら以下のコマンドでインストールすることができます。

sudo install apt-get install ~/Downloads/code_1.68.1-1655263094_amd64.deb

インストールしたら、「拡張機能」で次の機能をインストールします。

  • Docker
  • Japanese Language Pack for Visual Studio Code
  • Python
  • Remote-Containers

6. Docker composeのインストールと設定

6-1. Docker composeのインストール

こちらのページを参照し、Docker composeをインストールします。

6-2. Docker composeの設定

tensorflowの作業用フォルダを作成して、そのフォルダに.devcontainerのフォルダを作成して次の3つのファイルを作成します。

.devcontainer/
├── Dockerfile
├── devcontainer.json
└── docker-compose.yml

それぞれのファイルには以下の内容を記述します。

Dockerfile

dockerとUbuntuのユーザID, Group IDを同じにして、ファイルのアクセスに支障が出ないようにしています。

また、matplotlib, pandas, scikit-leanを追加でインストールしています。他に必要なライブラリがあれば追加・変更を行ってください。

FROM tensorflow/tensorflow:latest-gpu-jupyter

RUN apt-get update
# ユーザーを作成
ARG UID=1000
RUN useradd -m -u ${UID} user

# 作成したユーザーに切り替える
# このユーザーはRUN, CMD, ENTRYPOINT, docker run, exec の実行ユーザ。
USER ${UID}

# プロジェクトフォルダを/codeに追加する。ユーザー権限で扱えるようchownオプションを使う。
# ADDの実行権者はrootなのでオプションが必要。
ADD  --chown=user:user . /work
# 作成したフォルダに移動し、パッケージリストをインストールする。
#WORKDIR /work

RUN pip install matplotlib
RUN pip install scikit-learn
RUN pip install pandas

devcontainer.json

workspaceFolderとして/tfを指定しています。これはtensorflowのサンプルが/tfに保存されているのと、jupyter notebookのパスが/tfになっているためです。

{
	"name": "Existing Dockerfile",
	"context": "..",
	"dockerFile": "./Dockerfile",
	"workspaceFolder": "/tf",
	"settings": { 
		"terminal.integrated.shell.linux": "/bin/bash",
		"python.pythonPath": "/usr/local/bin/python",
		"python.linting.enabled": true,
		"python.linting.pylintEnabled": true,
		"python.linting.pylintPath": "/usr/local/bin/pylint"
	},
	"appPort": [ 9000 ],
	"remoteUser": "user",
	"extensions": [
		"ms-python.python"
		]
}

docker-compose.yml

volumesのところで、.devcontainerのフォルダのひとつ上をdockerの/tf/workにマウントするように指定しています。

ulimitsのところでGPUのメモリの制限をなくしています。デフォルトだとメモリ不足によりアプリケーションが動かないことがあります。

deploy > resources > reservations > devices、およびenvironmentのところでGPUを全て使用するような設定をしています。

exportsportsのところで、dockerのポートをubuntuの8888番のポートで使用できるようにしています。

version: '3'
services:
  tensorflow:
    build: .
    volumes:
      - ../:/tf/work
    working_dir: /tf
    ulimits:
      memlock: -1
      stack: -1
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              capabilities: [ gpu ]
    tty: true
    expose:
      - "8888"
    ports:
      - "127.0.0.1:8888:8888"
    environment:
      - NVIDIA_VISIBLE_DEVICES=all
      - NVIDIA_DRIVER_CAPABILITIES=all

これらのファイルを作成したら、VS Codeで作業用フォルダを開き、.devcontainer/docker-compose.ymlを左クリックして、Compose Upのメニューを選択します。

画面左端のDockerのアイコンをクリックして、今作成したDocker composeのイメージにVS Codeをアタッチします。

新しく開いたVS Codeのターミナル部分でnvidia-smiのコマンドを実行して、nvidiaのGPUが認識されていればOKです。

7. 別PCからリモートアクセスするための設定

以下の作業はMACまたはLinuxで使用できます。Windowsは開発に利用していないため設定方法は知りません。

7-1. SSHの鍵情報を作成する

mkdir ~/.ssh
cd ~/.ssh
ssh-keygen

id_rsa.pub, id_rsaなどのファイルが作成されます。

7-2. SSHの鍵情報をUbuntuに登録する

作成した鍵情報をUbuntuにコピーします。

ssh-copy-id ${ubuntuのユーザ名}@${UbuntuのIPアドレス}

ssh ${ubuntuのユーザ名}@${UbuntuのIPアドレス}のコマンドでパスワードがなくてもログインできることを確認します。

~/.ssh/configに以下の情報を追記します

Host Ubuntu
  HostName {UbuntuのIPアドレス}
  User {ubuntuのユーザ名}
  IdentityFile ~/.ssh/id_rsa
	LocalForward   8888 localhost:8888
  Port 22

LocalForwardはJupyter notebootのポートを別PCからアクセスできるようにするための設定です。

設定は以上になります。

8. 別PCから利用してみる

利用するPCのVS CodeにRemote - SSHの拡張機能をインストールすると、左端のメニューにリモートエクスプローラが表示されて、SSHターゲット先程登録したLinuxマシンが表示されます。フォルダのマークをクリックするとターゲットのPCのファイルを編集できるようになります。

クリックすると別の画面が開くので、そこでもリモートエクスプローラのメニューを選択して、

新しく開いたVS Codeの画面で「フォルダーを開く」をクリックして、フォルダには「/tf」フォルダを選択すると、docker上の/tfフォルダを編集できます。

tfフォルダにはworkフォルダが作成されていて、Ubuntuのファイルを閲覧・編集できます。

Dockerのファイルはイメージを消すとなくなりますので、workフォルダにファイルを保存するようにしましょう。

また、作業用PCでブラウザを起動して、URLにlocalhost:8888を指定するとjupyter notebookが開きます。tokenが必要というようなメッセージが表示されます。

dockerのVS Codeに戻り、ターミナルで次のコマンドを入力してください。

jupyter notebook list

このようにtokenが表示されます。

token=4fa1c9b0c0dacd567c41d32c05fa3b7ceb46bbf1b98ecd1bの部分をURLに追記してブラウザのアドレスバーに入力すると、ブラウザから接続先のDockerのjupyter notebookにアクセスできます。

http://localhost:8888?token=4fa1c9b0c0dacd567c41d32c05fa3b7ceb46bbf1b98ecd1b

これで、普段作業しているPCから機械学習用PCにリモートアクセスして、Docker上のPythonやJupyter notebookにアクセスできるようになりました。

ここに来るまで、5日ほど右往左往してしまいました。本ブログの内容をご活用して頂ければ幸いです。