Using a GPU in Linux Container during Development on Windows Host for Deep Learning
Want to be able to use your NVIDIA Graphical Processing Unit (GPU) on your Windows host machine while in a Linux Docker container? Deploying Deep Learning (DL) or general machine learning models is becoming an increasing occurrence. Along with the rise of deploying these models in microservices in Docker containers on Kubernetes. Deep Learning has the capability of increased training/testing/validation performance when using an NVIDIA GPU over a CPU.
CUDA and the current setup described is only for NVIDIA GPUS. For AMD cards there are alternatives using OpenCL.
Introduction
Linux containers use WLS/WLS2 (WSL will refer to both) to run on Windows devices. Until November 2021, Windows has not hosted the capability for the WLS to access the NVIDIA GPU. This recent Windows update:
https://www.cnet.com/tech/services-and-software/windows-10-november-2021-update-available-today/
enables this capability.
Before this time it would have been possible, but only if you were taking part in the Windows Insider program, which uses Beta CUDA drivers. This recent update now has made it accessible to all Windows 10 users.
The NVIDIA GPU is accessed using CUDA. CUDA is a parallel computing platform and the API that provides access to your NVIDIA GPU. CUDA is developed by Nvidia.
This article covers how to set up your Windows device to allow for NVIDIA GPU access and is tested using Deep Learning models in Python.
Before you Start
First, update Windows Version
Ensure that your Windows machine is on a version late enough to enable WLS2 to use the GPU. On Windows 10 this is version 21H2. https://docs.microsoft.com/en-us/windows/release-health/status-windows-10-21h2 . Or can update to Windows 11.
Second, ensure you have WLS2 Installed
There are many articles on this:
Briefly: Run Command Prompt as an administrator and enter.
wsl --install
Conveniently this works on Windows 10 devices and higher.
Next, Ensure Docker is using the WSL2 Engine
Assuming you already have Docker installed on your Windows machine. https://docs.docker.com/desktop/windows/install/
Under Settings>General ensure “Use the WSL 2 based engine” is ticked.
In Settings>Resources> WSL INTEGRATION ensure “Enable integration with my default WSL distro” is ticked.
You can check which is your default Linux Distro in the Command Prompt by entering:
wsl -l -v
Ensure you use your Linux distro as WSL default. (e.g. Ubuntu-20.04)
Getting CUDA on WSL
Open your WSL console by entering:
wsl
Or
If you have Window Terminal installed, open a new window under Ubuntu.
Install the NVIDIA Container Toolkit
In your WSL run:
sudo apt-get update
sudo apt-get install -y nvidia-docker2to install the docker tool kit.
To check that we have installed our NVIDIA tool kit correctly, we run:nvidia-smi , which displays the version, your GPU information. Similar to the example below:
Restart Docker
Test by running a docker container that uses your GPU
By benchmarking your GPU performance:
docker run --gpus all nvcr.io/nvidia/k8s/cuda-sample:nbody nbody -gpu -benchmarkLaunching and training a Deep Learning Framework:
docker run --gpus all -it --shm-size=1g --ulimit memlock=-1 --ulimit stack=67108864 nvcr.io/nvidia/tensorflow:20.03-tf2-py3Or exploring the performance in a Jupyter notebook:
docker run -it --gpus all -p 8888:8888 tensorflow/tensorflow:latest-gpu-py3-jupyterFor more details of this part of the tutorial: https://docs.nvidia.com/cuda/wsl-user-guide/index.html
Dev Containers in VS Code
In this section of the article, we go through our method of setting up our dev containers to develop and Deep Learning model using VS Code. Assuming you can already use Docker dev containers in VS Code.
Dockerfile
We use the tensorflow/tensorflow:lastest-gpu as a base image. Ensuring that we also get updates.
### GPU VersionFROM tensorflow/tensorflow:latest-gpuCMD pythonRUN apt-get updateRUN apt-get -y install gitRUN pip install --upgrade pip##############COPY . /appWORKDIR /appRUN set -ex && \pip3 install -r .devcontainer/requirements.txt
devcontainer.json
In the dev container ensure there are additional arguments in the "runArgs" option that allows you to share your GPU. Ensure you replace<memory-allocation> it with the amount of memory you wish to allocate for this.
"runArgs": ["--gpus","all","--shm-size","<memory-allocation>gb"],
requirements.txt
These are the files you need to install in your dev container using python-pip.
pytorch-lightningsklearnlightning-bolts
main.py
We use the iris data set from sklearn and pytorch-lightning with pytorch lightning bolts to access the GPU to train the model. The code below is an example of how to train such a model:
from pytorch_lightning import seed_everything, Trainerfrom pl_bolts.models.regression.logistic_regression import LogisticRegressionfrom pl_bolts.datamodules.sklearn_datamodule import SklearnDataModulefrom sklearn.datasets import load_irisseed_everything(1234)
# modelmodel = LogisticRegression(input_dim=4, num_classes=3, l1_strength=0.01, learning_rate=0.01)# dataX, y = load_iris(return_X_y=True)loaders = SklearnDataModule(X, y, num_workers=4)# traintrainer = Trainer(gpus=1,max_epochs=5,precision=16,auto_scale_batch_size=True,auto_select_gpus=True,auto_lr_find=True,)# tunetrainer.tune(model, datamodule=loaders)# fittrainer.fit(model, datamodule=loaders)
The current pytorch lightning bolts do not allow for viewing the results of the model. A fix has been requested:
GPU Utilisation
Below is an example of the Task Manager of my machine when using a different Deep Learning model, utilizing all of the GPU. Ensure that on the drop-down menus of the GPU performance graphics that Cuda is selected, to visualize the performance.
Thank you for reading, I hope it was somewhat helpful to you.
Any comments and suggestions are more than welcome.
