はじめに

注意

この投稿に記載されているは2021年2月21日時点での情報です。
Windows, Docker, WSL, グラフィックスドライバーなどのアップデートによりこの投稿の情報は使えなくなる可能性があります。
情報が古くなっていればコメントで教えてください。

TensorFlowの環境を作りたいけどWindowsのプロダクトキーを腐らせておくのもなんだかもったいない・・・。
そんな経緯からWindows上にTensorFlowの環境を作ることを検討したのですが、Windows内が散らかるのでやりたくありません。
そんなときにこちらの記事（待ってました CUDA on WSL 2）を見つけてうまくいったので現状この方法でいけますよという手順の紹介記事を書くことにしました。

手順

1. Windows Insider Preview Dev Channelをインストール

Windows Insider Programに参加してインストールまで済ませます。
自分はここで更新プログラムのビルドに失敗したのですが、クリーンインストールしてやりなおしたらいけました。
Microsoft公式の手順のページはこれ。Microsoft：Windows Insider Program を始める
画像とかがある個人の記事のほうがわかりやすいかもしれません。

2. グラフィックスドライバーをインストール

NVIDIA DEVELOPERのページ NVIDIA Developer：NVIDIA Drivers for CUDA on WSL, including DirectML Support からインストーラをダウンロードします。
ダウンロードにはアカウントが必要なのでなければ作ってください。
インストール時にはドライバーだけインストールするように指定できます。ドライバー以外をインストールされたくないなら指定してください。

3. WSLをインストール

Microsoft Storeから必要なディストリビューションのものをインストールします。コマンドからインストールすると最新版ではありませんでした。

WSLの更新もしておきます。

wsl --update

4. Docker Desktopをインストール

こちらのページ Docker Hub：Docker Desktop for Windows からDocker Desktop for Windowsのインストーラをダウンロードしてきてインストールします。
詳細な手順はこちら。 Docker ドキュメント日本語化プロジェクト：Windows に Docker Desktop をインストール

5. Docker for WindowsのWSL Integrationの有効化

Docker Desktopを起動して右上のアイコンの設定 > Resources > WSL Integration と進みます。
Enable integration with my default WSL distro のチェックボックスは初期状態でONになっています。
Enable integration with additional distros: の下にあるトグルスイッチをONにしてWSL Integratioを有効化します。

6. 動作確認

GPU対応のコンテナを起動してGPUを認識してくれていれば環境構築完了です。お疲れさまでした。
自分はTensorFlow公式のコンテナを起動してlist_local_devices()を実行してみました。

root@70f5ac4063c1:/tf# python
Python 3.6.9 (default, Oct  8 2020, 12:12:24)
[GCC 8.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> from tensorflow.python.client import device_lib
2021-02-22 11:37:02.874678: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcudart.so.11.0
>>> device_lib.list_local_devices()
2021-02-22 11:37:10.076496: I tensorflow/core/platform/cpu_feature_guard.cc:142] 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.
2021-02-22 11:37:10.077394: I tensorflow/compiler/jit/xla_gpu_device.cc:99] Not creating XLA devices, tf_xla_enable_xla_devices not set
2021-02-22 11:37:10.078752: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcuda.so.1
2021-02-22 11:37:10.273145: E tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:927] could not open file to read NUMA node: /sys/bus/pci/devices/0000:26:00.0/numa_node
Your kernel may have been built without NUMA support.
2021-02-22 11:37:10.273372: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1720] Found device 0 with properties:
pciBusID: 0000:26:00.0 name: GeForce RTX 3070 computeCapability: 8.6
coreClock: 1.755GHz coreCount: 46 deviceMemorySize: 8.00GiB deviceMemoryBandwidth: 417.29GiB/s
2021-02-22 11:37:10.273410: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcudart.so.11.0
2021-02-22 11:37:10.275010: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcublas.so.11
2021-02-22 11:37:10.275052: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcublasLt.so.11
2021-02-22 11:37:10.275620: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcufft.so.10
2021-02-22 11:37:10.275776: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcurand.so.10
2021-02-22 11:37:10.277315: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcusolver.so.10
2021-02-22 11:37:10.277615: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcusparse.so.11
2021-02-22 11:37:10.277728: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcudnn.so.8
2021-02-22 11:37:10.278004: E tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:927] could not open file to read NUMA node: /sys/bus/pci/devices/0000:26:00.0/numa_node
Your kernel may have been built without NUMA support.
2021-02-22 11:37:10.278351: E tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:927] could not open file to read NUMA node: /sys/bus/pci/devices/0000:26:00.0/numa_node
Your kernel may have been built without NUMA support.
2021-02-22 11:37:10.278538: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1862] Adding visible gpu devices: 0
2021-02-22 11:37:10.278585: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcudart.so.11.0
2021-02-22 11:37:10.747423: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1261] Device interconnect StreamExecutor with strength 1 edge matrix:
2021-02-22 11:37:10.747467: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1267]      0
2021-02-22 11:37:10.747473: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1280] 0:   N
2021-02-22 11:37:10.748282: E tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:927] could not open file to read NUMA node: /sys/bus/pci/devices/0000:26:00.0/numa_node
Your kernel may have been built without NUMA support.
2021-02-22 11:37:10.748463: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1489] Could not identify NUMA node of platform GPU id 0, defaulting to 0.  Your kernel may not have been built with NUMA support.
2021-02-22 11:37:10.748809: E tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:927] could not open file to read NUMA node: /sys/bus/pci/devices/0000:26:00.0/numa_node
Your kernel may have been built without NUMA support.
2021-02-22 11:37:10.749168: E tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:927] could not open file to read NUMA node: /sys/bus/pci/devices/0000:26:00.0/numa_node
Your kernel may have been built without NUMA support.
2021-02-22 11:37:10.749364: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1406] Created TensorFlow device (/device:GPU:0 with 6561 MB memory) -> physical GPU (device: 0, name: GeForce RTX 3070, pci bus id: 0000:26:00.0, compute capability: 8.6)
[name: "/device:CPU:0"
device_type: "CPU"
memory_limit: 268435456
locality {
}
incarnation: 2769908669862041284
, name: "/device:GPU:0"
device_type: "GPU"
memory_limit: 6880471296
locality {
  bus_id: 1
  links {
  }
}
incarnation: 9858467187116880750
physical_device_desc: "device: 0, name: GeForce RTX 3070, pci bus id: 0000:26:00.0, compute capability: 8.6"

投稿時点でのバージョン

Windows

OSビルド 1021318.1000

PS C:\WINDOWS\system32> wsl --list --verbose
  NAME                   STATE           VERSION
* Ubuntu-16.04           Running         2
  docker-desktop-data    Running         2
  Ubuntu-20.04           Running         2
  docker-desktop         Running         2

PS C:\WINDOWS\system32> wsl uname -r
5.4.91-microsoft-standard-WSL2

Docker

PS C:\WINDOWS\system32> docker --version
Docker version 20.10.2, build 2291f61

グラフィックスドライバー

NVIDIA グラフィックスドライバー 465.42

関連ページ

Qiita

待ってました CUDA on WSL 2

Microsoft

Windows Insider Program を始める

NVIDIA Developer

NVIDIA Drivers for CUDA on WSL, including DirectML Support

Docker Hub

Docker Desktop for Windows

Docker ドキュメント日本語化プロジェクト

Windows に Docker Desktop をインストール