wsl配置tensorflow环境

问题描述

查阅官网的 tensorflow 版本对照发现Windows 本机的 GPU 支持仅适用于 2.10 或更早版本，从 TF 2.11 开始，Windows 不支持 CUDA 构建。要在 Windows 上使用 TensorFlow GPU，您需要在 WSL2 中构建/安装 TensorFlow，或使用 TensorFlow-DirectML-Plugin 来构建/安装 tensorflow-cpu。

WSL 开发环境配置

WSL 安装

网上教程很多，大致流程差不多：

1. 打开控制面板->程序和功能->启用或关闭 windows 功能，把 适用于 Linux 的 Windows 子系统 和 Hyper-V 勾上
2. cmd 运行 bcdedit /set hypervisorlaunchtype auto
3. Micorsoft Store 安装 Ubuntu

Vscode 配置

下载插件 Remote Development，左上角选择要连接的 WSL

1. 创建工程文件夹 /home/username/project/
2. View->terminal 打开终端，运行以下两条命名：
   sudo apt-get update
sudo apt-get upgrade
3. 安装需要的 python 版本和 pip 工具
4. 创建并启用虚拟环境：
   pyton -m venv venv
source ./venv/bin/acitvate

安装 tensorflow

新版 tf 不用像以前一样找那么多插件的版本对应关系自己安装了。
打开 cmd 用 nvidia-smi 查看驱动版本

+---------------------------------------------------------------------------------------+
| NVIDIA-SMI 536.25                 Driver Version: 536.25       CUDA Version: 12.2     |
|-----------------------------------------+----------------------+----------------------+

结果显示我最高能用 cuda12.2。看官网介绍，cuda 12.2 对应 tensorflow 2.15.0。
直接运行 pip install tensorflow[and-gpu]==2.15.0 会报 Could not find TensorRT，手动安装 tensorrt 依然不行。在网上找到解决方案，于是执行：
pip install tensorflow[and-gpu]==2.15.0.post1
执行 python 代码：

import os
import tensorrt
import tensorflow as tf

print(tf.config.list_physical_devices('GPU'))

输出：

2024-04-18 17:48:13.893219: I tensorflow/core/util/port.cc:113] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
2024-04-18 17:48:13.918310: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
2024-04-18 17:48:13.918359: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
2024-04-18 17:48:13.919028: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
2024-04-18 17:48:13.922816: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: AVX2 AVX_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
2024-04-18 17:48:14.379458: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
True
2024-04-18 17:48:14.958324: I external/local_xla/xla/stream_executor/cuda/cuda_executor.cc:887] could not open file to read NUMA node: /sys/bus/pci/devices/0000:01:00.0/numa_node
Your kernel may have been built without NUMA support.
2024-04-18 17:48:14.995475: I external/local_xla/xla/stream_executor/cuda/cuda_executor.cc:887] could not open file to read NUMA node: /sys/bus/pci/devices/0000:01:00.0/numa_node
Your kernel may have been built without NUMA support.
2024-04-18 17:48:14.995533: I external/local_xla/xla/stream_executor/cuda/cuda_executor.cc:887] could not open file to read NUMA node: /sys/bus/pci/devices/0000:01:00.0/numa_node
Your kernel may have been built without NUMA support.
[PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')]

报了很多警告，我不知道怎么解决，但是从最后一行来看运行正常，不管了。
跑一个简单的 demo 试试：

import os
import tensorrt as trt
import tensorflow as tf
import keras

mnist = keras.datasets.mnist

(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0

model = keras.models.Sequential([
  keras.layers.Flatten(input_shape=(28, 28)),
  keras.layers.Dense(128, activation='relu'),
  keras.layers.Dropout(0.2),
  keras.layers.Dense(10)
])
predictions = model(x_train[:1]).numpy()
predictions
tf.nn.softmax(predictions).numpy()
loss_fn = keras.losses.SparseCategoricalCrossentropy(from_logits=True)
loss_fn(y_train[:1], predictions).numpy()
model.compile(optimizer='adam',
              loss=loss_fn,
              metrics=['accuracy'])
model.fit(x_train, y_train, epochs=5)
model.evaluate(x_test,  y_test, verbose=2)
probability_model = keras.Sequential([
  model,
  keras.layers.Softmax()
])
probability_model(x_test[:5])

结果：

2024-04-18 20:10:50.811251: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1929] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 3365 MB memory:  -> device: 0, name: NVIDIA GeForce RTX 4050 Laptop GPU, pci bus id: 0000:01:00.0, compute capability: 8.9
2024-04-18 20:10:58.280850: I external/local_tsl/tsl/platform/default/subprocess.cc:304] Start cannot spawn child process: No such file or directory
Epoch 1/5
2024-04-18 20:11:03.357516: I external/local_xla/xla/service/service.cc:168] XLA service 0x7f1d6c0b3390 initialized for platform CUDA (this does not guarantee that XLA will be used). Devices:
2024-04-18 20:11:03.357549: I external/local_xla/xla/service/service.cc:176]   StreamExecutor device (0): NVIDIA GeForce RTX 4050 Laptop GPU, Compute Capability 8.9
2024-04-18 20:11:03.384483: I tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.cc:269] disabling MLIR crash reproducer, set env var `MLIR_CRASH_REPRODUCER_DIRECTORY` to enable.
2024-04-18 20:11:03.870491: I external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:454] Loaded cuDNN version 8904
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
I0000 00:00:1713442263.964940  151868 device_compiler.h:186] Compiled cluster using XLA!  This line is logged at most once for the lifetime of the process.
1875/1875 [==============================] - 6s 2ms/step - loss: 0.2998 - accuracy: 0.9129
Epoch 2/5
1875/1875 [==============================] - 4s 2ms/step - loss: 0.1465 - accuracy: 0.9564
Epoch 3/5
1875/1875 [==============================] - 4s 2ms/step - loss: 0.1099 - accuracy: 0.9669
Epoch 4/5
1875/1875 [==============================] - 4s 2ms/step - loss: 0.0898 - accuracy: 0.9722
Epoch 5/5
1875/1875 [==============================] - 4s 2ms/step - loss: 0.0764 - accuracy: 0.9765
313/313 - 1s - loss: 0.0774 - accuracy: 0.9763 - 615ms/epoch - 2ms/step