Improve Cross-Platform Compatibility and Build Process

Dockerfile.cpu

@@ -0,0 +1,18 @@
+FROM --platform=linux/amd64 ubuntu:20.04
+
+RUN apt-get update && apt-get install -y libjpeg-dev libpng-dev build-essential git python3 python3-pip
+
+RUN pip3 install --upgrade pip
+RUN pip3 install torch
+RUN pip3 install aiofiles aiohttp auditwheel av defusedxml docutils gdown h5py ipython lmdb matplotlib numpy pandas Pillow pycocotools pytest pytorch_sphinx_theme Requests scipy setuptools Sphinx sphinx_gallery tabulate torch torchdata tqdm
+
+ENV FORCE_CUDA="0"
+ENV TORCH_CUDA_ARCH_LIST=""
+ENV USE_CUDA="0"
+
+COPY . /workspace
+WORKDIR /workspace
+RUN chmod +x /workspace/setup.py
+RUN python3 /workspace/setup.py develop
+
+CMD ["tail", "-f", "/dev/null"]

Dockerfile.gpu

@@ -0,0 +1,17 @@
+FROM pytorch/pytorch:1.7.0-cuda11.0-cudnn8-devel
+
+RUN apt-get update && apt-get install -y libjpeg-dev libpng-dev build-essential git
+
+RUN pip3 install --upgrade pip
+RUN pip3 install aiofiles aiohttp auditwheel av defusedxml docutils gdown h5py ipython lmdb matplotlib numpy pandas Pillow pycocotools pytest pytorch_sphinx_theme Requests scipy setuptools Sphinx sphinx_gallery tabulate torchdata tqdm
+
+ENV FORCE_CUDA="1"
+ENV TORCH_CUDA_ARCH_LIST="3.5 5.2 6.0 6.1 7.0+PTX"
+ENV TORCH_NVCC_FLAGS="-Xfatbin -compress-all"
+
+COPY . /workspace
+WORKDIR /workspace
+RUN chmod +x /workspace/setup.py
+RUN python3 /workspace/setup.py develop
+
+CMD ["tail", "-f", "/dev/null"]

docker-compose.yml

@@ -0,0 +1,44 @@
+# for mac arm chips see:
+# - https://docs.docker.com/desktop/troubleshoot/known-issues/
+
+# for nvidia gpus see:
+# - https://github.com/NVIDIA/nvidia-container-toolkit
+# - https://docs.docker.com/compose/gpu-support/
+
+# compiling pip dependencies:
+# 
+#   $ pipreqs . --mode no-pin --encoding utf-8
+#   $ mv requirements.txt requirements.in
+#   $ pip-compile requirements.in -o requirements.txt -vvv
+
+# usage: cpu
+# 
+#   $ docker-compose build
+#   $ docker-compose up
+
+# usage: gpu
+# 
+#   $ BUILD_TYPE=gpu docker-compose build
+#   $ BUILD_TYPE=gpu docker-compose up
+
+services:
+    main:
+      container_name: main
+      volumes:
+        - type: bind
+          source: .
+          target: /workspace
+      working_dir: /workspace
+      build:
+        context: .
+        dockerfile: Dockerfile.${DOCKER_FILE_SUFFIX:-cpu}
+      deploy:
+        resources:
+          reservations:
+            devices:
+              - driver: nvidia
+                count: all
+                capabilities: [gpu]
+      environment:
+        - NVIDIA_VISIBLE_DEVICES=${NVIDIA_VISIBLE_DEVICES:-all}
+        - NVIDIA_DRIVER_CAPABILITIES=${NVIDIA_DRIVER_CAPABILITIES:-compute,utility}

docker.sh

@@ -0,0 +1,12 @@
+#!/bin/bash
+
+if command -v nvidia-smi &> /dev/null && nvidia-smi &> /dev/null
+then
+    echo "NVIDIA GPU detected"
+    export DOCKER_FILE_SUFFIX=gpu
+else
+    echo "No NVIDIA GPU detected"
+    export DOCKER_FILE_SUFFIX=cpu
+fi
+
+docker-compose up --build