Merge branch 'gpus' of https://github.com/jmduarte/finn into jmduarte-gpus

docs/finn/getting_started.rst

@@ -133,7 +133,8 @@ These are summarized below:
 * (optional) ``IMAGENET_VAL_PATH`` specifies the path to the ImageNet validation directory for tests.
 * (optional) ``FINN_DOCKER_PREBUILT`` (default 0) if set to 1 then skip Docker image building and use the image tagged with ``FINN_DOCKER_TAG``.
 * (optional) ``FINN_DOCKER_RUN_AS_ROOT`` (default 0) if set to 1 then run Docker container as root, default is the current user.
-* (optional) ``FINN_DOCKER_GPU`` (autodetected) if not 0 then expose all Nvidia GPUs to Docker container for accelerated DNN training.
+* (optional) ``FINN_DOCKER_GPU`` (autodetected) if not 0 then expose all Nvidia GPUs or those selected by ``NVIDIA_VISIBLE_DEVICES`` to Docker container for accelerated DNN training. Requires `Nvidia Container Toolkit <https://github.com/NVIDIA/nvidia-docker>`_
+* (optional) ``NVIDIA_VISIBLE_DEVICES`` (default "") specifies specific Nvidia GPUs to use in Docker container. Possible values are a comma-separated list of GPU UUID(s) or index(es) e.g. ``0,1,2``, ``all``, ``none``, or void/empty/unset.
 
 Supported Hardware
 ===================

run-docker.sh

@@ -96,7 +96,8 @@ SCRIPTPATH=$(dirname "$SCRIPT")
 : ${FINN_DOCKER_TAG="xilinx/finn:$(git describe --tags --dirty).$XRT_DEB_VERSION"}
 : ${FINN_DOCKER_PREBUILT="0"}
 : ${FINN_DOCKER_RUN_AS_ROOT="0"}
-: ${FINN_DOCKER_GPU="docker info | grep nvidia | wc -m"}
+: ${FINN_DOCKER_GPU="$(docker info | grep nvidia | wc -m)"}
+: ${NVIDIA_VISIBLE_DEVICES=""}
 
 DOCKER_INTERACTIVE=""
 DOCKER_EXTRA=""
@@ -141,7 +142,11 @@ fi
 
 if [ "$FINN_DOCKER_GPU" != 0 ];then
   gecho "nvidia-docker detected, enabling GPUs"
-  DOCKER_EXTRA+="--gpus all"
+  if [ ! -z "$NVIDIA_VISIBLE_DEVICES" ];then
+    DOCKER_EXTRA+="--runtime nvidia -e NVIDIA_VISIBLE_DEVICES=$NVIDIA_VISIBLE_DEVICES "
+  else
+    DOCKER_EXTRA+="--gpus all"
+  fi
 fi
 
 VIVADO_HLS_LOCAL=$VIVADO_PATH