Commit 90c226cf authored by Ouail's avatar Ouail Committed by Francisco Massa

add the option to use a `FORCE_CUDA` to force cuda installation on docker (#612)

* add a FORCE_CUDA flag

Following discussion [here](https://github.com/facebookresearch/maskrcnn-benchmark/issues/167), this seemed the best solution

* Update Dockerfile

* Update setup.py

* add FORCE_CUDA as an ARG

* 	modified:   docker/Dockerfile
	modified:   setup.py

* small fix to readme of demo

* remove test print

* keep ARG_CUDA

* remove env value and use the one from ARG

* keep same formatting as source

* change proposed by @miguelvr

* Update INSTALL.md
parent 05feadff
...@@ -58,13 +58,17 @@ unset INSTALL_DIR ...@@ -58,13 +58,17 @@ unset INSTALL_DIR
### Option 2: Docker Image (Requires CUDA, Linux only) ### Option 2: Docker Image (Requires CUDA, Linux only)
Build image with defaults (`CUDA=9.0`, `CUDNN=7`): Build image with defaults (`CUDA=9.0`, `CUDNN=7`, `FORCE_CUDA=1`):
nvidia-docker build -t maskrcnn-benchmark docker/ nvidia-docker build -t maskrcnn-benchmark docker/
Build image with other CUDA and CUDNN versions: Build image with other CUDA and CUDNN versions:
nvidia-docker build -t maskrcnn-benchmark --build-arg CUDA=9.2 --build-arg CUDNN=7 docker/ nvidia-docker build -t maskrcnn-benchmark --build-arg CUDA=9.2 --build-arg CUDNN=7 docker/
Build image with FORCE_CUDA disabled:
nvidia-docker build -t maskrcnn-benchmark --build-arg FORCE_CUDA=0 docker/
Build and run image with built-in jupyter notebook(note that the password is used to log in jupyter notebook): Build and run image with built-in jupyter notebook(note that the password is used to log in jupyter notebook):
......
...@@ -38,7 +38,8 @@ docker run --rm -it \ ...@@ -38,7 +38,8 @@ docker run --rm -it \
-v /tmp/.X11-unix:/tmp/.X11-unix \ -v /tmp/.X11-unix:/tmp/.X11-unix \
--device=/dev/video0:/dev/video0 \ --device=/dev/video0:/dev/video0 \
--ipc=host maskrcnn-benchmark \ --ipc=host maskrcnn-benchmark \
python demo/webcam.py --min-image-size 300 python demo/webcam.py --min-image-size 300 \
--config-file configs/caffe2/e2e_mask_rcnn_R_50_FPN_1x_caffe2.yaml
``` ```
**DISCLAIMER:** *This was tested for an Ubuntu 16.04 machine, **DISCLAIMER:** *This was tested for an Ubuntu 16.04 machine,
......
...@@ -47,6 +47,8 @@ RUN git clone https://github.com/cocodataset/cocoapi.git \ ...@@ -47,6 +47,8 @@ RUN git clone https://github.com/cocodataset/cocoapi.git \
&& python setup.py build_ext install && python setup.py build_ext install
# install PyTorch Detection # install PyTorch Detection
ARG FORCE_CUDA="1"
ENV FORCE_CUDA=${FORCE_CUDA}
RUN git clone https://github.com/facebookresearch/maskrcnn-benchmark.git \ RUN git clone https://github.com/facebookresearch/maskrcnn-benchmark.git \
&& cd maskrcnn-benchmark \ && cd maskrcnn-benchmark \
&& python setup.py build develop && python setup.py build develop
......
...@@ -28,7 +28,7 @@ def get_extensions(): ...@@ -28,7 +28,7 @@ def get_extensions():
extra_compile_args = {"cxx": []} extra_compile_args = {"cxx": []}
define_macros = [] define_macros = []
if torch.cuda.is_available() and CUDA_HOME is not None: if (torch.cuda.is_available() and CUDA_HOME is not None) or os.getenv("FORCE_CUDA", "0") == "1":
extension = CUDAExtension extension = CUDAExtension
sources += source_cuda sources += source_cuda
define_macros += [("WITH_CUDA", None)] define_macros += [("WITH_CUDA", None)]
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment