Merge branch 'master' into kwargs_py_api

Author: narendasan

docker/Dockerfile

@@ -0,0 +1,49 @@
+ARG BASE=21.06
+ARG BASE_IMG=nvcr.io/nvidia/pytorch:${BASE}-py3
+FROM ${BASE_IMG} as base
+
+FROM base as trtorch-builder-base
+
+# Removing any bazel or trtorch pre-installed from the base image
+
+RUN rm -rf /opt/pytorch/trtorch /usr/bin/bazel
+
+RUN apt-get update && apt-get install --no-install-recommends -y curl gnupg
+RUN curl -fsSL https://bazel.build/bazel-release.pub.gpg | gpg --dearmor > /etc/apt/trusted.gpg.d/bazel.gpg
+RUN echo "deb [arch=amd64] https://storage.googleapis.com/bazel-apt stable jdk1.8" | tee /etc/apt/sources.list.d/bazel.list
+
+RUN apt-get update && apt-get install -y --no-install-recommends bazel-4.0.0
+RUN ln -s /usr/bin/bazel-4.0.0 /usr/bin/bazel
+
+# Workaround for bazel expecting both static and shared versions, we only use shared libraries inside container
+RUN cp /usr/lib/x86_64-linux-gnu/libnvinfer.so /usr/lib/x86_64-linux-gnu/libnvinfer_static.a
+
+RUN apt-get update && apt-get install -y --no-install-recommends locales ninja-build && rm -rf /var/lib/apt/lists/* && locale-gen en_US.UTF-8
+
+FROM trtorch-builder-base as trtorch-builder
+
+COPY . /workspace/trtorch/src
+WORKDIR /workspace/trtorch/src
+RUN cp ./docker/WORKSPACE.cu.docker  WORKSPACE
+
+# This script builds both libtrtorch bin/lib/include tarball and the Pythin wheel, in dist/
+RUN ./docker/dist-build.sh
+
+FROM base as trtorch
+
+# copy source repo
+COPY . /workspace/trtorch
+COPY --from=trtorch-builder  /workspace/trtorch/src/dist/ .
+RUN patch -u /opt/conda/lib/python3.8/site-packages/pytorch_quantization/nn/modules/tensor_quantizer.py -i /workspace/trtorch/docker/qat.patch
+RUN conda init bash
+
+RUN pip3 install ipywidgets --trusted-host pypi.org --trusted-host pypi.python.org --trusted-host=files.pythonhosted.org
+RUN jupyter nbextension enable --py widgetsnbextension
+
+RUN mkdir -p /opt/trtorch && tar xvf libtrtorch.tar.gz --strip-components 2 -C /opt/trtorch --exclude=LICENSE && pip3 install *.whl && rm -fr /workspace/trtorch/dist/*
+
+ENV LD_LIBRARY_PATH /opt/conda/lib/python3.8/site-packages/torch/lib:/opt/trtorch/lib:${LD_LIBRARY_PATH}
+ENV PATH /opt/trtorch/bin:${PATH}
+
+WORKDIR /workspace/trtorch/
+CMD /bin/bash

docker/Dockerfile.21.03

@@ -29,13 +29,25 @@ git_repository(
     shallow_since = "1570114335 -0400"
 )
 
+# External dependency for trtorch if you already have precompiled binaries.
+# This is currently used in pytorch NGC container CI testing.
+local_repository(
+    name = "trtorch",
+    path = "/opt/conda/lib/python3.8/site-packages/trtorch"
+)
+
 # CUDA should be installed on the system locally
 new_local_repository(
     name = "cuda",
     path = "/usr/local/cuda",
     build_file = "@//third_party/cuda:BUILD",
 )
 
+new_local_repository(
+    name = "cublas",
+    build_file = "@//third_party/cublas:BUILD",
+    path = "/usr",
+)
 
 ####################################################################################
 # Locally installed dependencies (use in cases of custom dependencies or aarch64)

docker/Dockerfile.21.06

@@ -0,0 +1,21 @@
+#!/bin/bash
+
+mkdir -p dist
+
+bazel build //:libtrtorch --compilation_mode opt
+
+cd py && MAX_JOBS=1 LANG=en_US.UTF-8 LANGUAGE=en_US:en LC_ALL=en_US.UTF-8 python3 setup.py bdist_wheel --use-cxx11-abi
+
+cd ..
+
+cp bazel-bin/libtrtorch.tar.gz dist/
+cp py/dist/* dist/
+
+pip3 install ipywidgets --trusted-host pypi.org --trusted-host pypi.python.org --trusted-host=files.pythonhosted.org
+jupyter nbextension enable --py widgetsnbextension
+
+pip3 install timm
+
+# test install
+mkdir -p /opt/trtorch && tar xvf dist/libtrtorch.tar.gz --strip-components 2 -C /opt/trtorch --exclude=LICENSE && pip3 uninstall -y trtorch && pip3 install dist/*.whl
+

docker/Dockerfile.21.07

@@ -0,0 +1,8 @@
+#!/bin/bash
+
+pip3 install timm --trusted-host pypi.org --trusted-host pypi.python.org --trusted-host=files.pythonhosted.org
+# Build and run unit tests
+cd tests/modules && python3 ./hub.py
+cd ../..
+
+bazel test //tests:tests //tests:python_api_tests --compilation_mode=opt --jobs=4

docker/WORKSPACE.cu.docker

@@ -0,0 +1,19 @@
+--- torch/nn/functional.py	2021-10-01 16:53:42.827338664 -0700
++++ functional.py	2021-10-01 16:53:34.639338618 -0700
+@@ -4975,7 +4975,7 @@
+         f"was expecting embedding dimension of {embed_dim_to_check}, but got {embed_dim}"
+     if isinstance(embed_dim, torch.Tensor):
+         # embed_dim can be a tensor when JIT tracing
+-        head_dim = embed_dim.div(num_heads, rounding_mode='trunc')
++        head_dim = int(embed_dim.div(num_heads, rounding_mode='trunc'))
+     else:
+         head_dim = embed_dim // num_heads
+     assert head_dim * num_heads == embed_dim, f"embed_dim {embed_dim} not divisible by num_heads {num_heads}"
+@@ -5044,6 +5044,7 @@
+     #
+     # reshape q, k, v for multihead attention and make em batch first
+     #
++    bsz = int(bsz)
+     q = q.contiguous().view(tgt_len, bsz * num_heads, head_dim).transpose(0, 1)
+     if static_k is None:
+         k = k.contiguous().view(k.shape[0], bsz * num_heads, head_dim).transpose(0, 1)

docker/dist-build.sh

@@ -0,0 +1,11 @@
+--- /opt/conda/lib/python3.8/site-packages/pytorch_quantization/nn/modules/tensor_quantizer.py	2021-08-16 22:50:37.000000000 +0000
++++ tensor_quantizer.py	2021-10-19 20:41:54.288077426 +0000
+@@ -291,7 +291,7 @@
+             quant_dim = list(amax.shape).index(list(amax_sequeeze.shape)[0])
+             scale = amax_sequeeze / bound
+             outputs = torch.fake_quantize_per_channel_affine(
+-                inputs, scale.data, torch.zeros_like(scale, dtype=torch.long).data, quant_dim,
++                inputs, scale.data, torch.zeros_like(scale, dtype=torch.int32).data, quant_dim,
+                 -bound - 1 if not self._unsigned else 0, bound)
+ 
+         return outputs

docker/dist-test.sh

@@ -234,7 +234,7 @@ def run(self):
       long_description=long_description,
       ext_modules=ext_modules,
       install_requires=[
-          'torch>=1.9.0+cu111,<1.10.0',
+          'torch>=1.9.0<1.11.0',
       ],
       setup_requires=[],
       cmdclass={