Tensorflow GPU with Docker in WSL2 using VS Code Dev Containers
- Windows 11
- Windows Subsystem for Linux (WSL2)
- Latest NVIDIA Windows GPU Driver
- Docker Desktop WSL2
- Visual Studio Code with Dev Containers extension
- Create a new repository from the template or clone the code in WSL
- In
Dockerfile
:- Update the TensorFlow Docker image: run
nvidia-smi
command in your Terminal, and get the CUDA version. Then find the Tensorflow version compatible with the CUDA version here - Update
USERNAME
ARG to create a non-root user - Update
USER_UID
andUSER_GID
ARGs (useid
command in your Terminal to get them) - Update required packages (optional)
- Update the TensorFlow Docker image: run
- Update pip packages in
requirements.txt
(optional) - Open your project in WSL with VS Code and reopen in Container: (F1) use the Dev Containers: Reopen in Container command
Run pyhton main.py
to test it.