Use Tensorflow and PyTorch on WSL 2
Tensorflow 2.10 has been the last version that supports CUDA on Windows, so where should the Wnidows users go with Tensorflow GPU? The answer is WSL 2. You may have heard the performance issues of WSL 2 but that would not be much of an issue for CUDA. By Apr 2023, PyTorch 2.0 (without compile) and Tensorflow 2.12 can run perfectly well on Ubuntu WSL 2 without losing much performance. Here are the installation steps:
- Install Conda
- Install WSL 2
- Install CUDA following the CUDA on WSL User Guide
- Install Tensorflow following the official guide. Do not install
tensorflow-gpu
, which does not support WSL. - (Optional) Install TensorRT
DirectML is a different way of using CUDA from Windows native. By Apr 2023, DirectML has about 10-20% performance loss. Better use CUDA directly in WSL.
If you install TensorRT via pip, update your $CONDA_PREFIX/etc/conda/activate.d/env_vars.sh
as follows:
TENSORRT_PATH=$(dirname $(python -c "import tensorrt;print(tensorrt.__file__)"))
CUDNN_PATH=$(dirname $(python -c "import nvidia.cudnn;print(nvidia.cudnn.__file__)"))
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$CONDA_PREFIX/lib/:$CUDNN_PATH/lib:$TENSORRT_PATH
Test GPU:
python -c "import tensorflow as tf; print(tf.config.list_physical_devices('GPU'))"
Important Note: In WSL reading files from Windows drives can seriously affect IO performance. I strongly suggest you move your code and data file into the WSL drive. In WSL, copy files from C drive by cp /mnt/c/...
.
To move WSL virtual disk to another drive, follow this blog
After you export and import WSL, you might need to update the default user in WSL. Follow this link.