diff --git a/Dockerfile b/Dockerfile index e5bd901a33c..2e4007b8e9c 100644 --- a/Dockerfile +++ b/Dockerfile @@ -50,10 +50,12 @@ COPY . . RUN git submodule update --init --recursive FROM conda as build +ARG CMAKE_VARS WORKDIR /opt/pytorch COPY --from=conda /opt/conda /opt/conda COPY --from=submodule-update /opt/pytorch /opt/pytorch RUN --mount=type=cache,target=/opt/ccache \ + export eval ${CMAKE_VARS} && \ TORCH_CUDA_ARCH_LIST="3.5 5.2 6.0 6.1 7.0+PTX 8.0" TORCH_NVCC_FLAGS="-Xfatbin -compress-all" \ CMAKE_PREFIX_PATH="$(dirname $(which conda))/../" \ python setup.py install diff --git a/README.md b/README.md index 30c26139818..61c1ee5e9ff 100644 --- a/README.md +++ b/README.md @@ -362,11 +362,19 @@ should increase shared memory size either with `--ipc=host` or `--shm-size` comm The `Dockerfile` is supplied to build images with CUDA 11.1 support and cuDNN v8. You can pass `PYTHON_VERSION=x.y` make variable to specify which Python version is to be used by Miniconda, or leave it unset to use the default. + ```bash make -f docker.Makefile # images are tagged as docker.io/${your_docker_username}/pytorch ``` +You can also pass the `CMAKE_VARS="..."` environment variable to specify additional CMake variables to be passed to CMake during the build. +See [setup.py](./setup.py) for the list of available variables. + +```bash +CMAKE_VARS="BUILD_CAFFE2=ON BUILD_CAFFE2_OPS=ON" make -f docker.Makefile +``` + ### Building the Documentation To build documentation in various formats, you will need [Sphinx](http://www.sphinx-doc.org) and the diff --git a/docker.Makefile b/docker.Makefile index fd49964c458..5c77711b68c 100644 --- a/docker.Makefile +++ b/docker.Makefile @@ -12,6 +12,7 @@ CUDA_VERSION = 11.7.0 CUDNN_VERSION = 8 BASE_RUNTIME = ubuntu:18.04 BASE_DEVEL = nvidia/cuda:$(CUDA_VERSION)-cudnn$(CUDNN_VERSION)-devel-ubuntu18.04 +CMAKE_VARS ?= # The conda channel to use to install cudatoolkit CUDA_CHANNEL = nvidia @@ -31,7 +32,8 @@ BUILD_ARGS = --build-arg BASE_IMAGE=$(BASE_IMAGE) \ --build-arg CUDA_CHANNEL=$(CUDA_CHANNEL) \ --build-arg PYTORCH_VERSION=$(PYTORCH_VERSION) \ --build-arg INSTALL_CHANNEL=$(INSTALL_CHANNEL) \ - --build-arg TRITON_VERSION=$(TRITON_VERSION) + --build-arg TRITON_VERSION=$(TRITON_VERSION) \ + --build-arg CMAKE_VARS="$(CMAKE_VARS)" EXTRA_DOCKER_BUILD_FLAGS ?= BUILD ?= build