[WIP] Barebone VPF + PyTorch extension docker image build (#294) (#295)

* [WIP] Barebone VPF + PyTorch extension docker image build (#294)

* virtual env build

* Cmake list updates

* vpf + torch working

* newline

* formatting

* Barebone VPF (+ PyTorch extension)  (#296)

* [#293] Barebone VPF docker build

* Newline

Co-authored-by: Harsimrat Sandhawalia <10599550+sandhawalia@users.noreply.github.com>

README.md

@@ -42,14 +42,26 @@ cp -a $(VIDEO_CODEC_SDK) Video_Codec_SDK
 ```
 docker-compose -f docker/docker-compose.yml build vpf
 # Get test sample
-wget http://www.scikit-video.org/stable/_static/bikes.mp4 $HOME/Downloads/
+wget http://www.scikit-video.org/stable/_static/bikes.mp4 -P $HOME/Downloads/
 # run image
 docker-compose -f docker/docker-compose.yml run -v $HOME/Downloads:/Downloads vpf
 # or this way
 docker run  -it --gpus=all -e NVIDIA_DRIVER_CAPABILITIES=video,compute,utility -v $HOME/Downloads:/Downloads nvidia/videoprocessingframework:vpf
 python Tests.py 0 /Downloads/bikes.mp4 /Downloads/bikes-vpf.mp4
 ```
 
+5. Build & Run image with `pytorch` extension
+
+```
+docker-compose -f docker/docker-compose.yml build --build-arg GEN_PYTORCH_EXT=1 vpf
+# Get test sample
+wget http://www.scikit-video.org/stable/_static/bikes.mp4 -P $HOME/Downloads/
+# run image
+docker-compose -f docker/docker-compose.yml run -v $HOME/Downloads:/Downloads vpf
+# Run predictions on video
+python SampleTorchResnet.py 0 /Downloads/bikes.mp4
+```
+
 You can build [`tensorrt`](https://developer.nvidia.com/tensorrt) enabled image by replacing `vpf` with `vpf-tensorrt` in the above steps and test the following.
 
 ```

docker/Dockerfile-tensorrt.gpu

@@ -15,7 +15,7 @@ RUN apt update && DEBIAN_FRONTEND=noninteractive \
     libavcodec-dev libavformat-dev libswscale-dev \
     libxvidcore-dev x264 libx264-dev libfaac-dev \
     libmp3lame-dev libtheora-dev libvorbis-dev \
-    python3.8 python3-pip python3-virtualenv && \
+    python3.8 python3-pip python3.8-venv && \
     rm -rf /var/lib/apt/lists/*
 
 # https://python-poetry.org/docs/
@@ -36,7 +36,8 @@ COPY . .
 
 RUN make -f docker/Makefile vpf_built \
                             PYTHON_BINARY=$PYTHON_BINARY \
-                            VIDEO_CODEC_SDK=/opt/Video_Codec_SDK
+                            VIDEO_CODEC_SDK=/opt/Video_Codec_SDK \
+                            GEN_PYTORCH_EXT=1
 
 FROM nvcr.io/nvidia/tensorrt:21.08-py3
 
@@ -49,7 +50,7 @@ RUN apt-get update && \
     libdc1394-22-dev libssl-dev \
     libxine2-dev libv4l-dev \
     libboost-all-dev libdc1394-22-dev \
-    python3.8 python3-pip python3-virtualenv curl
+    python3.8 python3-pip python3.8-venv curl
 
 WORKDIR /opt
 

docker/Dockerfile.gpu

@@ -8,14 +8,16 @@ ENV POETRY_HOME="/opt/poetry" \
 ENV PATH="/opt/.venv/bin:$POETRY_HOME/bin:$PATH" \
     PYTHON_BINARY="/opt/.venv/bin/python"
 
+ARG GEN_PYTORCH_EXT
+
 RUN apt update && DEBIAN_FRONTEND=noninteractive \
     apt -y install \
     git wget cmake \
     build-essential curl \
     libavcodec-dev libavformat-dev libswscale-dev \
     libxvidcore-dev x264 libx264-dev libfaac-dev \
     libmp3lame-dev libtheora-dev libvorbis-dev \
-    python3.8 python3-pip python3-virtualenv && \
+    python3.8 python3-pip python3.8-venv && \
     rm -rf /var/lib/apt/lists/*
 
 # https://python-poetry.org/docs/
@@ -29,14 +31,15 @@ COPY poetry.lock poetry.lock
 COPY pyproject.toml pyproject.toml
 
 # Build poetry virtual environment
-RUN make -f docker/Makefile build_env EXTRAS=vpf
+RUN if [ "$GEN_PYTORCH_EXT" = "1" ] ; then make -f docker/Makefile build_env EXTRAS=vpf-pytorch ; else make -f docker/Makefile build_env EXTRAS=vpf ; fi
 
 # Build vpf
 COPY . .
 
 RUN make -f docker/Makefile vpf_built \
                             PYTHON_BINARY=$PYTHON_BINARY \
-                            VIDEO_CODEC_SDK=/opt/Video_Codec_SDK
+                            VIDEO_CODEC_SDK=/opt/Video_Codec_SDK \
+                            GEN_PYTORCH_EXT=$GEN_PYTORCH_EXT
 
 FROM docker.io/nvidia/cuda:11.2.0-cudnn8-runtime-ubuntu20.04
 
@@ -49,7 +52,7 @@ RUN apt-get update && \
     libdc1394-22-dev libssl-dev \
     libxine2-dev libv4l-dev \
     libboost-all-dev libdc1394-22-dev \
-    python3.8 python3-pip python3-virtualenv curl
+    python3.8 python3-pip python3.8-venv curl
 
 WORKDIR /opt
 
@@ -63,7 +66,7 @@ COPY . .
 ENV POETRY_HOME="/opt/poetry"
 
 ENV PATH="/opt/.venv/bin:$POETRY_HOME/bin:$PATH" \
-    LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/opt/dist/bin:/opt/.venv/lib/python3.8/site-packages/torch/lib" \
-    PYTHONPATH=/opt/.venv/bin:/opt/dist/bin
+    LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/opt/dist:/opt/dist/bin:/opt/.venv/lib/python3.8/site-packages/torch/lib" \
+    PYTHONPATH=/opt/.venv/bin:/opt/dist:/opt/dist/bin
 
 ENTRYPOINT ["/bin/bash"]

docker/Makefile

@@ -11,9 +11,9 @@ vpf_built:
 	  -D NVCUVID_LIBRARY=$(VIDEO_CODEC_SDK)/Lib/linux/stubs/x86_64/libnvcuvid.so \
 	  -D NVENCODE_LIBRARY=$(VIDEO_CODEC_SDK)/Lib/linux/stubs/x86_64/libnvidia-encode.so \
 	  -D CMAKE_INSTALL_PREFIX=../dist \
-	  -D GENERATE_PYTORCH_EXTENSION=on \
+	  -D GENERATE_PYTORCH_EXTENSION:BOOL="$(GEN_PYTORCH_EXT)" \
 	  -D PYTHON_BINARY=$(PYTHON_BINARY) \
-	  -D GENERATE_PYTHON_BINDINGS=ON .. && \
+	  -D GENERATE_PYTHON_BINDINGS:BOOL="1" .. && \
   make -j$(nproc) && make install
 
 linter_python: $(shell find . -type f -name "*.py" -not -path "*venv*")

docker/docker-compose.yml

@@ -6,6 +6,8 @@ services:
     build:
       context: ../
       dockerfile: docker/Dockerfile.gpu
+      args:
+        GEN_PYTORCH_EXT: 0
     ipc: host
     runtime: nvidia
     image: nvidia/videoprocessingframework:vpf