passing core tests ?

Signed-off-by: Oleksii Kuchaiev <okuchaiev@nvidia.com>

Jenkinsfile

@@ -24,9 +24,9 @@ pipeline {
         sh 'python setup.py style'
       }
     }
-    stage('Unittests general') {
+    stage('Unittests Core') {
       steps {
-        sh './reinstall.sh && python -m unittest tests/*.py'
+        sh './reinstall.sh && python -m unittest tests/core/*.py'
       }
     }
     stage('Unittests ASR') {

nemo/backends/pytorch/common/zero_data.py

@@ -18,11 +18,11 @@ def neuralType2TensorShape(neural_type: NeuralType, default_dim=32, skip_batch_a
       torch.Size
     """
     dims = []
-    for axis_ind, axis_type in neural_type.axis2type.items():
-        if axis_type._semantics == BatchTag and skip_batch_axis:
+    for axis in neural_type.axes:
+        if axis.kind == AxisKind.Batch and skip_batch_axis:
             continue
-        if axis_type.dim is not None:
-            dims.append(axis_type.dim)
+        if axis.size is not None:
+            dims.append(axis.size)
         else:
             dims.append(default_dim)
     return torch.Size(dims)

nemo/core/neural_types/axes.py

@@ -69,7 +69,7 @@ class AxisType(object):
            is_list (bool, default=False):
     """
 
-    def __init__(self, kind: AxisKindAbstract, size: Optional[int], is_list=False):
+    def __init__(self, kind: AxisKindAbstract, size: Optional[int] = None, is_list=False):
         if size is not None and is_list:
             raise ValueError("The axis can't be list and have a fixed size")
         self.kind = kind

nemo/core/neural_types/elements.py

@@ -28,6 +28,9 @@
     'LabelsType',
     'LossType',
     'RegressionValuesType',
+    'PredictionsType',
+    'LogprobsType',
+    'LengthsType',
 ]
 import abc
 from abc import ABC, abstractmethod
@@ -55,7 +58,6 @@ def fields(self) -> Optional[Tuple]:
 
     def compare(self, second) -> NeuralTypeComparisonResult:
         # First, check general compatibility
-        result = NeuralTypeComparisonResult.SAME
         first_t = type(self)
         second_t = type(second)
 
@@ -108,11 +110,21 @@ def __str__(self):
         return "neural type representing logits"
 
 
+class LogprobsType(ElementType):
+    def __str__(self):
+        return "neural type representing log probabilities"
+
+
 class LabelsType(ElementType):
     def __str__(self):
         return "neural type representing labels"
 
 
+class LengthsType(ElementType):
+    def __str__(self):
+        return "neural type representing lengths of something"
+
+
 class LossType(ElementType):
     def __str__(self):
         return "neural type representing loss value"
@@ -151,6 +163,11 @@ def __str__(self):
         return "mfcc spectorgram type"
 
 
-class RegressionValuesType(ElementType):
+class PredictionsType(ElementType):
+    def __str__(self):
+        return "predictions values type"
+
+
+class RegressionValuesType(PredictionsType):
     def __str__(self):
         return "regression values type"

nemo/core/neural_types/neural_type.py

@@ -44,7 +44,7 @@ def __init__(self, elements_type: ElementType = VoidType(), axes: Optional[Tuple
             )
         self.elements_type = elements_type
         if axes is not None:
-            self.__check_sanity(axes)
+            NeuralType.__check_sanity(axes)
             axes_list = []
             for axis in axes:
                 if isinstance(axis, str):
@@ -63,9 +63,7 @@ def compare(self, second) -> NeuralTypeComparisonResult:
         axes_a = self.axes
         axes_b = second.axes
 
-        kinds_a = dict()
-        kinds_b = dict()
-
+        # "Big void" type
         if isinstance(self.elements_type, VoidType) and self.axes is None:
             return NeuralTypeComparisonResult.SAME
 
@@ -75,28 +73,29 @@ def compare(self, second) -> NeuralTypeComparisonResult:
             else:
                 return NeuralTypeComparisonResult.INCOMPATIBLE
 
-        dimensions_pass = True
-        for axis_a, axis_b in zip(axes_a, axes_b):
-            kinds_a[axis_a.kind] = axis_a.size
-            kinds_b[axis_b.kind] = axis_b.size
-            if axis_a.kind != axis_b.kind or axis_a.is_list != axis_b.is_list:
-                dimensions_pass = False
-
-        if kinds_a.keys() != kinds_b.keys():
-            return NeuralTypeComparisonResult.INCOMPATIBLE
-        for kind, size in kinds_a.items():
-            if size != kinds_b[kind]:
-                return NeuralTypeComparisonResult.DIM_INCOMPATIBLE
-
+        dimensions_pass = NeuralType.__compare_axes(axes_a, axes_b)
         element_comparison_result = self.elements_type.compare(second.elements_type)
-        if dimensions_pass:
+
+        # SAME DIMS
+        if dimensions_pass == 0:
             return element_comparison_result
-        elif element_comparison_result == NeuralTypeComparisonResult.SAME:
-            return NeuralTypeComparisonResult.TRANSPOSE_SAME
+        # TRANSPOSE_SAME DIMS
+        elif dimensions_pass == 1:
+            if element_comparison_result == NeuralTypeComparisonResult.SAME:
+                return NeuralTypeComparisonResult.TRANSPOSE_SAME
+            else:
+                return NeuralTypeComparisonResult.INCOMPATIBLE
+        # DIM_INCOMPATIBLE DIMS
+        elif dimensions_pass == 2:
+            if element_comparison_result == NeuralTypeComparisonResult.SAME:
+                return NeuralTypeComparisonResult.DIM_INCOMPATIBLE
+            else:
+                return NeuralTypeComparisonResult.INCOMPATIBLE
         else:
             return NeuralTypeComparisonResult.INCOMPATIBLE
 
-    def __check_sanity(self, axes):
+    @staticmethod
+    def __check_sanity(axes):
         # check that list come before any tensor dimension
         are_strings = True
         for axis in axes:
@@ -119,6 +118,50 @@ def __check_sanity(self, axes):
                 "You have list dimension after Tensor dimension. All list dimensions must preceed Tensor dimensions"
             )
 
+    @staticmethod
+    def __compare_axes(axes_a, axes_b) -> int:
+        """
+        Compares axes_a and axes_b
+        Args:
+            axes_a: first axes tuple
+            axes_b: second axes tuple
+
+        Returns:
+            0 - if they are exactly the same
+            1 - if they are "TRANSPOSE_SAME"
+            2 - if the are "DIM_INCOMPATIBLE"
+            3 - if they are different
+        """
+        if axes_a is None and axes_b is None:
+            return 0
+        elif axes_a is None and axes_b is not None:
+            return 3
+        elif axes_a is not None and axes_b is None:
+            return 3
+        elif len(axes_a) != len(axes_b):
+            return 3
+        # After these ifs we know that len(axes_a) == len(axes_b)
+
+        same = True
+        kinds_a = dict()
+        kinds_b = dict()
+        for axis_a, axis_b in zip(axes_a, axes_b):
+            kinds_a[axis_a.kind] = axis_a.size
+            kinds_b[axis_b.kind] = axis_b.size
+            if axis_a.kind != axis_b.kind or axis_a.is_list != axis_b.is_list or axis_a.size != axis_b.size:
+                same = False
+        if same:
+            return 0
+        else:
+            # can be TRANSPOSE_SAME, DIM_INCOMPATIBLE
+            if kinds_a.keys() == kinds_b.keys():
+                for key, value in kinds_a.items():
+                    if kinds_b[key] != value:
+                        return 2
+                return 1
+            else:
+                return 3
+
 
 class NmTensor(NeuralType):
     """Class representing data which flows between NeuralModules' ports.

requirements/requirements_test.txt

@@ -4,3 +4,4 @@ pytest-runner
 black
 isort[requirements]
 wrapt
+onnxruntime

tests/core/test_deploy_export.py

@@ -0,0 +1,149 @@
+# ! /usr/bin/python
+# -*- coding: utf-8 -*-
+
+# Copyright 2019 NVIDIA. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# =============================================================================
+
+# import os
+# from pathlib import Path
+#
+# # git clone git@github.com:microsoft/onnxruntime.git
+# # cd onnxruntime
+# # ./build.sh --update --build --config RelWithDebInfo --build_shared_lib --parallel --use_cuda \
+# #            --cudnn_home /usr/lib/x86_64-linux-gnu --cuda_home /usr/local/cuda --enable_pybind --build_wheel
+# # pip install --upgrade ./build/Linux/RelWithDebInfo/dist/onnxruntime_gpu-1.1.0-cp37-cp37m-linux_x86_64.whl
+# import onnxruntime as ort
+# import torch
+# from ruamel.yaml import YAML
+#
+# import nemo
+# import nemo.collections.asr as nemo_asr
+# import nemo.collections.nlp as nemo_nlp
+# import nemo.collections.nlp.nm.trainables.common.token_classification_nm
+# from tests.common_setup import NeMoUnitTest
+#
+#
+# class TestDeployExport(NeMoUnitTest):
+#     def setUp(self):
+#         """ Setups neural factory so it will use GPU instead of CPU. """
+#         NeMoUnitTest.setUp(self)
+#
+#         # Perform computations on GPU.
+#         self.nf._placement = nemo.core.DeviceType.GPU
+#
+#     def __test_export_route(self, module, out_name, mode, input_example=None):
+#         out = Path(out_name)
+#         if out.exists():
+#             os.remove(out)
+#
+#         self.nf.deployment_export(module=module, output=out_name, input_example=input_example, d_format=mode)
+#
+#         self.assertTrue(out.exists())
+#         if mode == nemo.core.DeploymentFormat.ONNX:
+#             if isinstance(input_example, tuple):
+#                 outputs_fwd = module.forward(*input_example)
+#             else:
+#                 outputs_fwd = module.forward(input_example)
+#             sess_options = ort.SessionOptions()
+#             sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_EXTENDED
+#             ort_session = ort.InferenceSession(out_name, sess_options)
+#             inputs = dict()
+#             input_names = list(module.input_ports)
+#             for i in range(len(input_names)):
+#                 input_name = (
+#                     "encoded_lengths"
+#                     if type(module).__name__ == "JasperEncoder" and input_names[i] == "length"
+#                     else input_names[i]
+#                 )
+#                 inputs[input_name] = (
+#                     input_example[i].cpu().numpy() if isinstance(input_example, tuple) else input_example.cpu().numpy()
+#                 )
+#             outputs_ort = ort_session.run(None, inputs)
+#             outputs_ort = torch.from_numpy(outputs_ort[0]).cuda()
+#             self.assertLess(
+#                 (outputs_ort - (outputs_fwd[0] if isinstance(outputs_fwd, tuple) else outputs_fwd)).norm(p=2), 5.0e-4
+#             )
+#         if out.exists():
+#             os.remove(out)
+#
+#     def test_simple_module_export(self):
+#         simplest_module = nemo.backends.pytorch.tutorials.TaylorNet(dim=4)
+#         self.__test_export_route(
+#             module=simplest_module,
+#             out_name="simple.pt",
+#             mode=nemo.core.DeploymentFormat.TORCHSCRIPT,
+#             input_example=None,
+#         )
+#
+#     def test_TokenClassifier_module_export(self):
+#         t_class = nemo.collections.nlp.nm.trainables.common.token_classification_nm.TokenClassifier(
+#             hidden_size=512, num_classes=16, use_transformer_pretrained=False
+#         )
+#         self.__test_export_route(
+#             module=t_class,
+#             out_name="t_class.pt",
+#             mode=nemo.core.DeploymentFormat.TORCHSCRIPT,
+#             input_example=torch.randn(16, 16, 512).cuda(),
+#         )
+#
+#     def test_TokenClassifier_module_onnx_export(self):
+#         t_class = nemo.collections.nlp.nm.trainables.common.token_classification_nm.TokenClassifier(
+#             hidden_size=512, num_classes=16, use_transformer_pretrained=False
+#         )
+#         self.__test_export_route(
+#             module=t_class,
+#             out_name="t_class.onnx",
+#             mode=nemo.core.DeploymentFormat.ONNX,
+#             input_example=torch.randn(16, 16, 512).cuda(),
+#         )
+#
+#     def test_jasper_decoder_export_ts(self):
+#         j_decoder = nemo_asr.JasperDecoderForCTC(feat_in=1024, num_classes=33)
+#         self.__test_export_route(
+#             module=j_decoder, out_name="j_decoder.ts", mode=nemo.core.DeploymentFormat.TORCHSCRIPT, input_example=None
+#         )
+#
+#     def test_hf_bert_ts(self):
+#         bert = nemo.collections.nlp.nm.trainables.common.huggingface.BERT(pretrained_model_name="bert-base-uncased")
+#         input_example = (
+#             torch.randint(low=0, high=16, size=(2, 16)).cuda(),
+#             torch.randint(low=0, high=1, size=(2, 16)).cuda(),
+#             torch.randint(low=0, high=1, size=(2, 16)).cuda(),
+#         )
+#         self.__test_export_route(
+#             module=bert, out_name="bert.ts", mode=nemo.core.DeploymentFormat.TORCHSCRIPT, input_example=input_example
+#         )
+#
+#     def test_hf_bert_pt(self):
+#         bert = nemo.collections.nlp.nm.trainables.common.huggingface.BERT(pretrained_model_name="bert-base-uncased")
+#         self.__test_export_route(module=bert, out_name="bert.pt", mode=nemo.core.DeploymentFormat.PYTORCH)
+#
+#     def test_jasper_encoder_to_onnx(self):
+#         with open("tests/data/jasper_smaller.yaml") as file:
+#             yaml = YAML(typ="safe")
+#             jasper_model_definition = yaml.load(file)
+#
+#         jasper_encoder = nemo_asr.JasperEncoder(
+#             conv_mask=False,
+#             feat_in=jasper_model_definition['AudioToMelSpectrogramPreprocessor']['features'],
+#             **jasper_model_definition['JasperEncoder']
+#         )
+#
+#         self.__test_export_route(
+#             module=jasper_encoder,
+#             out_name="jasper_encoder.onnx",
+#             mode=nemo.core.DeploymentFormat.ONNX,
+#             input_example=(torch.randn(16, 64, 256).cuda(), torch.randn(256).cuda()),
+#         )