DeepFaceLab GPU-enabled Dockerfile [Ubuntu 16.04 w/ CUDA v9, CUDNN v7]

Dockerfile

# NOTE: nvidia-docker must be enabled on the Docker host

FROM nvidia/cuda:9.0-base-ubuntu16.04
ENV LANG C.UTF-8

# adapted from the official tensorflow docker builds [https://github.com/tensorflow/tensorflow/blob/master/tensorflow/tools/dockerfiles/dockerfiles/devel-gpu.Dockerfile]
RUN apt-get update && apt-get install -y --no-install-recommends \
        build-essential \
        cuda-command-line-tools-9-0 \
        cuda-cublas-9-0 \
        cuda-cufft-9-0 \
        cuda-curand-9-0 \
        cuda-cusolver-9-0 \
        cuda-cusparse-9-0 \
        libcudnn7=7.2.1.38-1+cuda9.0 \
        libnccl2=2.2.13-1+cuda9.0 \
        libfreetype6-dev \
        libhdf5-serial-dev \
        libpng12-dev \
        libzmq3-dev \
        pkg-config \
        software-properties-common \
        unzip \
        && \
    apt-get clean && \
    rm -rf /var/lib/apt/lists/*

RUN apt-get update && \
        apt-get install nvinfer-runtime-trt-repo-ubuntu1604-4.0.1-ga-cuda9.0 && \
        apt-get update && \
        apt-get install libnvinfer4=4.1.2-1+cuda9.0

RUN apt-get update && apt-get install -y \
    cmake \
    git \
    python3 \
    python3-pip \
    libsm6 \
    libxext6 \
    libxrender-dev

RUN pip3 install --upgrade \
    pip \
    setuptools

WORKDIR /workspace/DeepFaceLab

COPY requirements-gpu-cuda9-cudnn7.txt /workspace/DeepFaceLab

RUN pip3 install -r requirements-gpu-cuda9-cudnn7.txt

Will set up a Docker-ized environment for GPU-enabled sessions of DeepFaceLab functions. The NVIDIA drivers and nvidia-docker runtime must be installed on host machine for GPU support. This is designed to set up environment with the docker run part executing just one shell command inside mounted volumes.

Example:

docker run \ 
    --runtime=nvidia  \
    -v $(pwd):/workspace/DeepFaceLab \
    -v $(pwd)/faceA:/workspace/DeepFaceLab/faceA \
    -v $(pwd)/faceB:/workspace/DeepFaceLab/faceB \
    -v $(pwd)/model:/workspace/DeepFaceLab/model \
    -it deepfacelab \
    bash -c "python3 main.py train --training-data-src-dir faceA --training-data-dst-dir faceB/aligned --model-dir model --model LIAEF128"

This can be converted to initiate a Jupyter Notebooks to work in, instead of running shell command directly. See example integration here.