Export parameters propagated, requirements enforced. (#1096)

* Propagating onnx.export() parameters; requiring torch>=1.6

* Adding all requirements install for Dockerfile and reinstall.sh, moving to 20.08 base

* Bumping megatron-lm version

Signed-off-by: Boris Fomitchev <[email protected]>

Co-authored-by: Oleksii Kuchaiev <[email protected]>

Dockerfile

@@ -14,7 +14,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-ARG BASE_IMAGE=nvcr.io/nvidia/pytorch:20.07-py3
+ARG BASE_IMAGE=nvcr.io/nvidia/pytorch:20.08-py3
+
 
 # build an image that includes only the nemo dependencies, ensures that dependencies
 # are included first for optimal caching, and useful for building a development
@@ -32,8 +33,8 @@ RUN apt-get update && \
 
 # install nemo dependencies
 WORKDIR /tmp/nemo
-COPY requirements/requirements_docker.txt requirements.txt
-RUN pip install --disable-pip-version-check --no-cache-dir -r requirements.txt
+COPY requirements .
+RUN for f in requirements/*.txt; do pip install --disable-pip-version-check --no-cache-dir -r $f; done
 
 # copy nemo source into a scratch image
 FROM scratch as nemo-src

nemo/core/classes/exportable.py

@@ -51,6 +51,10 @@ def export(
         output: str,
         input_example=None,
         output_example=None,
+        verbose=False,
+        export_params=True,
+        do_constant_folding=True,
+        keep_initializers_as_inputs=False,
         onnx_opset_version: int = 12,
         try_script: bool = False,
         set_eval: bool = True,
@@ -132,20 +136,22 @@ def export(
                             _out_example = self.forward(*_in_example)
                         else:
                             _out_example = self.forward(_in_example)
+
                     torch.onnx.export(
                         jitted_model,
                         _in_example,
                         output,
                         input_names=input_names,
                         output_names=output_names,
-                        verbose=False,
-                        export_params=True,
-                        do_constant_folding=True,
-                        keep_initializers_as_inputs=True,
+                        verbose=verbose,
+                        export_params=export_params,
+                        do_constant_folding=do_constant_folding,
+                        keep_initializers_as_inputs=keep_initializers_as_inputs,
                         dynamic_axes=dynamic_axes,
                         opset_version=onnx_opset_version,
                         example_outputs=_out_example,
                     )
+
                     # Verify the model can be read, and is valid
                     onnx_model = onnx.load(output)
                     onnx.checker.check_model(onnx_model, full_check=True)

reinstall.sh

@@ -1,16 +1,22 @@
 #!/usr/bin/env bash
 set -e
 
+PIP=pip
+
 echo 'Uninstalling stuff'
-pip uninstall -y nemo_toolkit
+${PIP} uninstall -y nemo_toolkit
 
 # Kept for legacy purposes
-pip uninstall -y nemo_asr
-pip uninstall -y nemo_nlp
-pip uninstall -y nemo_tts
-pip uninstall -y nemo_simple_gan
+${PIP} uninstall -y nemo_asr
+${PIP} uninstall -y nemo_nlp
+${PIP} uninstall -y nemo_tts
+${PIP} uninstall -y nemo_simple_gan
+
+${PIP} install -U setuptools
+
+for f in requirements/*.txt; do ${PIP} install ${PIP_FLAGS}--disable-pip-version-check --no-cache-dir -r $f; done 
 
 echo 'Installing stuff'
-pip install -e ".[all]"
+${PIP} install -e ".[all]"
 
 echo 'All done!'

requirements/requirements_nlp.txt

@@ -9,5 +9,5 @@ numpy
 tqdm>=4.41.0
 rapidfuzz
 gdown
-megatron-lm>=1.1.3
+megatron-lm>=1.1.4
 inflect

setup.py

@@ -89,7 +89,6 @@ def req_file(filename, folder="requirements"):
 
 extras_require = {
     # User packages
-    'docker': req_file("requirements_docker.txt"),
     'test': req_file("requirements_test.txt"),
     # Collections Packages
     'asr': req_file("requirements_asr.txt"),