目的

GPU+tensorflow環境をwsl2+dockerで実現する。なるべくクリーンな環境で試したい。

環境

windows11
wsl2(ubuntu)
NVIDIA GeForce GTX 1050Ti

手順

1. windows11にNVIDIA Driverをインストールする。

僕の場合、勝手にインストールされていた気がします。インストールされていない場合、https://www.nvidia.co.jp/Download/index.aspx?lang=jp でダウンロードできます。
ubuntu上でnvidia-smiを実行し、GPUを認識できていれば成功です。

$ nvidia-smi
Sun Mar 20 11:00:03 2022
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 510.47.03    Driver Version: 511.65       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 ...  On   | 00000000:1C:00.0  On |                  N/A |
| 44%   31C    P0    N/A /  75W |   1744MiB /  4096MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|  No running processes found                                                 |
+-----------------------------------------------------------------------------+

wsl2はwsl1と違い、本物のlinux kernelを使用しているのでこの時点でlinuxがGPUを認識できていると考えて大丈夫です。

2. dockerをwsl2上にインストールする。

ここからはnvidia公式サイト(https://docs.nvidia.com/cuda/wsl-user-guide/index.html#getting-started-with-cuda-on-wsl) の4.2節を流用しているだけです。
wsl2のターミナル上で以下のコマンドを実行します。docker公式でlinuxへのインストール手順があるのでこれを参考にします。(https://docs.docker.jp/linux/step_one.html)

curl -fsSL https://get.docker.com/ | sh

成功すると以下のようなログが出力されます。

curl -fsSL https://get.docker.com/ | sh
# Executing docker install script, commit: 93d2499759296ac1f9c510605fef85052a2c32be

WSL DETECTED: We recommend using Docker Desktop for Windows.
Please get Docker Desktop from https://www.docker.com/products/docker-desktop


You may press Ctrl+C now to abort this script.
+ sleep 20
+ sudo -E sh -c apt-get update -qq >/dev/null
[sudo] password for USERNAME:
+ sudo -E sh -c DEBIAN_FRONTEND=noninteractive apt-get install -y -qq apt-transport-https ca-certificates curl >/dev/null
+ sudo -E sh -c curl -fsSL "https://download.docker.com/linux/ubuntu/gpg" | gpg --dearmor --yes -o /usr/share/keyrings/docker-archive-keyring.gpg
+ sudo -E sh -c echo "deb [arch=amd64 signed-by=/usr/share/keyrings/docker-archive-keyring.gpg] https://download.docker.com/linux/ubuntu focal stable" > /etc/apt/sources.list.d/docker.list
+ sudo -E sh -c apt-get update -qq >/dev/null
+ sudo -E sh -c DEBIAN_FRONTEND=noninteractive apt-get install -y -qq --no-install-recommends  docker-ce-cli docker-scan-plugin docker-ce >/dev/null
+ version_gte 20.10
+ [ -z  ]
+ return 0
+ sudo -E sh -c DEBIAN_FRONTEND=noninteractive apt-get install -y -qq docker-ce-rootless-extras >/dev/null

================================================================================

To run Docker as a non-privileged user, consider setting up the
Docker daemon in rootless mode for your user:

    dockerd-rootless-setuptool.sh install

Visit https://docs.docker.com/go/rootless/ to learn about rootless mode.


To run the Docker daemon as a fully privileged service, but granting non-root
users access, refer to https://docs.docker.com/go/daemon-access/

WARNING: Access to the remote API on a privileged Docker daemon is equivalent
         to root access on the host. Refer to the 'Docker daemon attack surface'
         documentation for details: https://docs.docker.com/go/attack-surface/

================================================================================

wslなのでdocker desktopをお勧めすると書いてありますが、docker desktopは有償化されてしまったので、無視します。(個人の利用では無料)
dockerとターミナル上に打ち込み、コマンド認識できていればOKです。

3. NVIDIA Container Toolkitをwsl2上にインストールする。

以下のコマンドでtoolkitをインストールします。何をやっているか正直理解できていないですが、公式なので大丈夫でしょう。

distribution=$(. /etc/os-release;echo $ID$VERSION_ID)

curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add -

curl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.list | sudo tee /etc/apt/sources.list.d/nvidia-docker.list

インストールをします。

sudo apt-get update

sudo apt-get install -y nvidia-docker2

4. 動作確認

dockerがGPUを認識できているか動作確認を行います。
まずdockerを再起動しておきましょう。

sudo service docker stop

sudo service docker start

docker imageをプルしてきて、テストします。

$ sudo docker run --rm --gpus all nvcr.io/nvidia/k8s/cuda-sample:nbody nbody -gpu -benchmark
Unable to find image 'nvcr.io/nvidia/k8s/cuda-sample:nbody' locally
nbody: Pulling from nvidia/k8s/cuda-sample
2f94e549220a: Pull complete
8a196a8ba405: Pull complete
d84024e05c66: Pull complete
c6e912b1cab0: Pull complete
2c392aca5e2c: Pull complete
e12411a32402: Pull complete
246511fe8354: Pull complete
8eec16b7bc5c: Pull complete
9917b9983432: Pull complete
Digest: sha256:2a17caedda57b64f02f85f6ffae682ebbfdeadbf401e1c6a4c561f1b12b2257a
Status: Downloaded newer image for nvcr.io/nvidia/k8s/cuda-sample:nbody
Run "nbody -benchmark [-numbodies=<numBodies>]" to measure performance.
        -fullscreen       (run n-body simulation in fullscreen mode)
        -fp64             (use double precision floating point values for simulation)
        -hostmem          (stores simulation data in host memory)
        -benchmark        (run benchmark to measure performance)
        -numbodies=<N>    (number of bodies (>= 1) to run in simulation)
        -device=<d>       (where d=0,1,2.... for the CUDA device to use)
        -numdevices=<i>   (where i=(number of CUDA devices > 0) to use for simulation)
        -compare          (compares simulation results running once on the default GPU and once on the CPU)
        -cpu              (run n-body simulation on the CPU)
        -tipsy=<file.bin> (load a tipsy model file for simulation)

NOTE: The CUDA Samples are not meant for performance measurements. Results may vary when GPU Boost is enabled.

> Windowed mode
> Simulation data stored in video memory
> Single precision floating point simulation
> 1 Devices used for simulation
GPU Device 0: "Pascal" with compute capability 6.1

> Compute 6.1 CUDA device: [NVIDIA GeForce GTX 1050 Ti]
6144 bodies, total time for 10 iterations: 5.805 ms
= 65.027 billion interactions per second
= 1300.547 single-precision GFLOP/s at 20 flops per interaction

認識できていますね。

5. tensorflowのdocker imageをプルしてテスト

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

これでbashに入れます。bash上でpythonを実行して、確認します。

Python 3.8.10 (default, Nov 26 2021, 20:14:08)
[GCC 9.3.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>>
>>>
>>> from tensorflow.python.client import device_lib
lib.list_local_devices()
>>> device_lib.list_local_devices()
2022-03-20 02:33:15.693566: I tensorflow/core/platform/cpu_feature_guard.cc:151] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 FMA
To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.
2022-03-20 02:33:15.827200: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:922] could not open file to read NUMA node: /sys/bus/pci/devices/0000:1c:00.0/numa_node
Your kernel may have been built without NUMA support.
2022-03-20 02:33:15.834763: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:922] could not open file to read NUMA node: /sys/bus/pci/devices/0000:1c:00.0/numa_node
Your kernel may have been built without NUMA support.
2022-03-20 02:33:15.835051: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:922] could not open file to read NUMA node: /sys/bus/pci/devices/0000:1c:00.0/numa_node
Your kernel may have been built without NUMA support.
2022-03-20 02:33:16.524473: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:922] could not open file to read NUMA node: /sys/bus/pci/devices/0000:1c:00.0/numa_node
Your kernel may have been built without NUMA support.
2022-03-20 02:33:16.524820: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:922] could not open file to read NUMA node: /sys/bus/pci/devices/0000:1c:00.0/numa_node
Your kernel may have been built without NUMA support.
2022-03-20 02:33:16.524865: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1609] Could not identify NUMA node of platform GPU id 0, defaulting to 0.  Your kernel may not have been built with NUMA support.
2022-03-20 02:33:16.525211: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:922] could not open file to read NUMA node: /sys/bus/pci/devices/0000:1c:00.0/numa_node
Your kernel may have been built without NUMA support.
2022-03-20 02:33:16.525316: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1525] Created device /device:GPU:0 with 2752 MB memory:  -> device: 0, name: NVIDIA GeForce GTX 1050 Ti, pci bus id: 0000:1c:00.0, compute capability: 6.1
[name: "/device:CPU:0"
device_type: "CPU"
memory_limit: 268435456
locality {
}
incarnation: 15340632351876512816
xla_global_id: -1
, name: "/device:GPU:0"
device_type: "GPU"
memory_limit: 2886126798
locality {
  bus_id: 1
  links {
  }
}
incarnation: 10473275661079893773
physical_device_desc: "device: 0, name: NVIDIA GeForce GTX 1050 Ti, pci bus id: 0000:1c:00.0, compute capability: 6.1"
xla_global_id: 416903419
]

GPUを認識できていますね。

まとめ

wsl2を使うと、docker desktopを使わなくてもdockerを使用できます。かなりクリーンな環境でGPUを利用できるので良いと思います。