TneosrRT 7.1 ONNX-Parser open source release

onnx · Jul 1, 2020 · 1cf89dd · 1cf89dd
1 parent 2eb74d9
commit 1cf89dd
Show file tree

Hide file tree

Showing 35 changed files with 1,908 additions and 984 deletions.
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -1,4 +1,4 @@
- # Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.
+ # Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
  #
  # Permission is hereby granted, free of charge, to any person obtaining a
  # copy of this software and associated documentation files (the "Software"),
@@ -42,7 +42,7 @@ set(PARSER_LINKER_SCRIPT  ${ONNX2TRT_ROOT}/libnvonnxparser.version)
 # Version information
 #--------------------------------------------------
 set(ONNX2TRT_MAJOR 7)
-set(ONNX2TRT_MINOR 0)
+set(ONNX2TRT_MINOR 1)
 set(ONNX2TRT_PATCH 0)
 
 #--------------------------------------------------
@@ -56,6 +56,8 @@ set(IMPORTER_SOURCES
   onnx2trt_utils.cpp
   ShapedWeights.cpp
   ShapeTensor.cpp
+  LoopHelpers.cpp
+  RNNHelpers.cpp
   OnnxAttrs.cpp
 )
 

diff --git a/ImporterContext.hpp b/ImporterContext.hpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.
+ * Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
  *
  * Permission is hereby granted, free of charge, to any person obtaining a
  * copy of this software and associated documentation files (the "Software"),
@@ -48,8 +48,10 @@ class ImporterContext final : public IImporterContext
         mTensorNameCounts; // Keep track of how many times a tensor name shows up, to avoid duplicate naming in TRT.
     StringMap<size_t>
         mLayerNameCounts; // Keep track of how many times a tensor name shows up, to avoid duplicate naming in TRT.
-    std::unordered_set<std::string> mUnsupportedShapeTensors; // Container to hold any shape tensors that are the output of layers that do not support shape tensors.
+    std::unordered_set<std::string> mUnsupportedShapeTensors; // Container to hold output tensor names of layers that produce shape tensor outputs but do not natively support them.
     StringMap<std::string> mLoopTensors; // Container to map subgraph tensors to their original outer graph names.
+    std::string mOnnxFileLocation; // Keep track of the directory of the parsed ONNX file
+
 public:
     ImporterContext(nvinfer1::INetworkDefinition* network, nvinfer1::ILogger* logger)
         : _network(network)
@@ -88,7 +90,14 @@ class ImporterContext final : public IImporterContext
     {
         return mLoopTensors;
     }
-
+    virtual void setOnnxFileLocation(std::string location) override
+    {
+        mOnnxFileLocation = location;
+    }
+    virtual std::string getOnnxFileLocation() override
+    {
+        return mOnnxFileLocation;
+    }
     // This actually handles weights as well, but is named this way to be consistent with the tensors()
     virtual void registerTensor(TensorOrWeights tensor, const std::string& basename) override
     {

diff --git a/LICENSE b/LICENSE
@@ -1,7 +1,7 @@
 MIT License
 
-Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.
-Copyright (c) 2018 Open Neural Network Exchange
+Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
+Copyright (c) 2020 Open Neural Network Exchange
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal

diff --git a/LoopHelpers.cpp b/LoopHelpers.cpp
@@ -0,0 +1,18 @@
+#include "LoopHelpers.hpp"
+#include "onnx2trt_utils.hpp"
+
+namespace onnx2trt
+{
+
+nvinfer1::ITensor* addLoopCounter(IImporterContext* ctx, nvinfer1::ILoop* loop, int32_t initial)
+{
+    nvinfer1::ITensor* initialTensor = addConstantScalar(ctx, initial, ::ONNX_NAMESPACE::TensorProto::INT32, nvinfer1::Dims{1, 1})->getOutput(0);
+    nvinfer1::ITensor* one = addConstantScalar(ctx, 1, ::ONNX_NAMESPACE::TensorProto::INT32, nvinfer1::Dims{1, 1})->getOutput(0);
+
+    auto counter = loop->addRecurrence(*initialTensor);
+    nvinfer1::ITensor* addOne = ctx->network()->addElementWise(*counter->getOutput(0), *one, nvinfer1::ElementWiseOperation::kSUM)->getOutput(0);
+    counter->setInput(1, *addOne);
+    return counter->getOutput(0);
+}
+
+} // namespace onnx2trt
diff --git a/LoopHelpers.hpp b/LoopHelpers.hpp
@@ -0,0 +1,12 @@
+#pragma once
+
+#include <NvInfer.h>
+
+#include "ImporterContext.hpp"
+
+namespace onnx2trt
+{
+
+nvinfer1::ITensor* addLoopCounter(IImporterContext* ctx, nvinfer1::ILoop* loop, int32_t initial = 0);
+
+} // namespace onnx2trt
diff --git a/ModelImporter.cpp b/ModelImporter.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.
+ * Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
  *
  * Permission is hereby granted, free of charge, to any person obtaining a
  * copy of this software and associated documentation files (the "Software"),
@@ -31,6 +31,7 @@
 #include <google/protobuf/text_format.h>
 
 #include <limits>
+#include <functional>
 #include <unordered_set>
 
 namespace onnx2trt
@@ -81,8 +82,7 @@ Status setStringMap(
     return Status::success();
 }
 
-Status parseGraph(
-    IImporterContext* ctx, const ::ONNX_NAMESPACE::GraphProto& graph, bool deserializingINetwork, int* currentNode)
+Status parseGraph(IImporterContext* ctx, const ::ONNX_NAMESPACE::GraphProto& graph, bool deserializingINetwork, int* currentNode)
 {
     // Import initializers.
     for (const ::ONNX_NAMESPACE::TensorProto& initializer : graph.initializer())
@@ -129,17 +129,19 @@ Status parseGraph(
         LOG_VERBOSE(ssInputs.str());
 
         // Dispatch to appropriate converter.
-        if (!opImporters.count(node.op_type()))
+        const NodeImporter* importFunc{nullptr};
+        if (opImporters.count(node.op_type()))
         {
-            return MAKE_ERROR("No importer registered for op: " + node.op_type(), ErrorCode::kUNSUPPORTED_NODE);
+            importFunc = &opImporters.at(node.op_type());
+        }
+        else
+        {
+            LOG_INFO("No importer registered for op: " << node.op_type() << ". Attempting to import as plugin.");
+            importFunc = &opImporters.at("FallbackPluginImporter");
         }
-        const NodeImporter& importFunc = opImporters.at(node.op_type());
         std::vector<TensorOrWeights> outputs;
 
-        const int nbLayersBefore = ctx->network()->getNbLayers();
-        GET_VALUE(importFunc(ctx, node, nodeInputs), &outputs);
-
-        ctx->registerLayer(ctx->network()->getLayer(nbLayersBefore), node.name());
+        GET_VALUE((*importFunc)(ctx, node, nodeInputs), &outputs);
 
         if (deserializingINetwork)
         {
@@ -333,6 +335,7 @@ bool ModelImporter::supportsModel(
         }
     }
     auto* ctx = &_importer_ctx;
+
     auto checkForInput = [&input_node, &ctx](::ONNX_NAMESPACE::NodeProto const& node) {
         for (auto input : node.input())
         {
@@ -344,6 +347,28 @@ bool ModelImporter::supportsModel(
         return false;
     };
 
+    auto checkForShapeTensors = [&ctx](::ONNX_NAMESPACE::NodeProto const& node){
+        for (int i = 0; i < ctx->network()->getNbInputs(); i++)
+        {
+            auto input = ctx->network()->getInput(i);
+            if (input->isShapeTensor())
+            {
+                if (input->getType() == nvinfer1::DataType::kFLOAT || node.op_type() == "Loop" || node.op_type() == "Scan")
+                {
+                    auto name = input->getName();
+                    for (auto input : node.input())
+                    {
+                        if (input == name)
+                        {
+                            return true;
+                        }
+                    }
+                }
+            }
+        }
+        return false;
+    };
+
     bool newSubGraph(true);
     // Sort and partition supported subgraphs
     std::vector<size_t> topological_order;
@@ -356,19 +381,19 @@ bool ModelImporter::supportsModel(
     for (int node_idx : topological_order)
     {
         ::ONNX_NAMESPACE::NodeProto const& node = model.graph().node(node_idx);
-
         // Add the node to the subgraph if:
-        //     1. Importer function is registered for the operator type
-        //     2. It is NOT directly connected to an unsupported input
-        //     3. Parsing did NOT hit an error on the node
-        //     4. Any shape tensor output is coming from a supported node
+        //     1. There is an importer function registered for the operator type
+        //     2. It is not directly connected to an unsupported input
+        //     3. It is not directly connected to an unsupported shape tensor input
+        //     4. It did not illegally produce a shape tensor output
+        //     5. It was sucessfully parsed
         bool registered = supportsOperator(node.op_type().c_str());
         bool containsInput = (input_node.empty()) ? false : checkForInput(node);
+        bool containsShapeInput = checkForShapeTensors(node);
+        auto tensorName = node.name();
+        bool supportedShapeTensor = ctx->unsupportedShapeTensors().count(tensorName) == 0 ? true : false;
         bool containsIndex = node_idx == error_node;
-        auto const tensor = node.output(0);
-        bool supportedShapeTensorOutput = ctx->unsupportedShapeTensors().count(tensor) == 0 ? true : false;
-
-        if (registered && !containsInput && !containsIndex && supportedShapeTensorOutput)
+        if (registered && !containsInput && !containsShapeInput && supportedShapeTensor && !containsIndex)
         {
             if (newSubGraph)
             {
@@ -383,6 +408,7 @@ bool ModelImporter::supportsModel(
         }
         else
         {
+            std::cout << "Found unsupported node: " << tensorName << std::endl;
             // This is not a supported node, reset newSubGraph
             newSubGraph = true;
             allSupported = false;
@@ -440,31 +466,29 @@ void removeShapeTensorCasts(IImporterContext* ctx)
     for (int i = 0, e = ctx->network()->getNbLayers(); i < e; ++i)
     {
         nvinfer1::ILayer* layer = ctx->network()->getLayer(i);
-        constexpr nvinfer1::DataType SHAPE_TENSOR_TYPE = nvinfer1::DataType::kINT32;
         if (layer->getNbOutputs() > 0 && layer->getOutput(0)->isShapeTensor())
         {
-            layer->resetPrecision();
             layer->resetOutputType(0);
-            layer->setPrecision(SHAPE_TENSOR_TYPE);
-            layer->setOutputType(0, SHAPE_TENSOR_TYPE);
             nvinfer1::ITensor& t = *layer->getOutput(0);
+            // Assume that boolean tensors were not cast, and thus have their type correctly set.
+            const nvinfer1::DataType shapeTensorType = t.getType() == nvinfer1::DataType::kBOOL ? nvinfer1::DataType::kBOOL : nvinfer1::DataType::kINT32;
+            layer->setOutputType(0, shapeTensorType);
             // Set type only if necessary, to avoid TensorRT warnings
             // about setting type of non-input/output tensors.
-            if (t.getType() != SHAPE_TENSOR_TYPE)
+            if (t.getType() != shapeTensorType)
             {
-                t.setType(SHAPE_TENSOR_TYPE);
+                t.setType(shapeTensorType);
             }
             // Some layers do not support shape tensor outputs. Keep track of these tensor names
             // for supportsModel().
             auto type = layer->getType();
-            auto elementwiseOp = layer->getType() == nvinfer1::LayerType::kELEMENTWISE ? (static_cast<nvinfer1::IElementWiseLayer*>(layer))->getOperation() : nvinfer1::ElementWiseOperation::kSUM;
-            auto reduceOp = layer->getType() == nvinfer1::LayerType::kREDUCE ? (static_cast<nvinfer1::IReduceLayer*>(layer))->getOperation() : nvinfer1::ReduceOperation::kSUM;
-
+            auto elementwiseOp = type == nvinfer1::LayerType::kELEMENTWISE ? (static_cast<nvinfer1::IElementWiseLayer*>(layer))->getOperation() : nvinfer1::ElementWiseOperation::kSUM;
+            auto reduceOp = type == nvinfer1::LayerType::kREDUCE ? (static_cast<nvinfer1::IReduceLayer*>(layer))->getOperation() : nvinfer1::ReduceOperation::kSUM;
             if (!supportsShapeTensor(type, elementwiseOp, reduceOp))
             {
-                auto name = layer->getOutput(0)->getName();
+                auto name = layer->getName();
                 ctx->unsupportedShapeTensors().insert(name);
-                LOG_ERROR("Found " << name << " as a shape tensor output from a layer that does not support it!");
+                LOG_ERROR("Found unsupported shape-tensor producing layer:" << name);
             }
         }
     }
@@ -477,7 +501,7 @@ Status ModelImporter::importModel(
     auto* ctx = &_importer_ctx;
     _importer_ctx.clearOpsets();
     // Initialize plugin registry
-    initLibNvInferPlugins(static_cast<void*>(&ctx->logger()), "ONNXTRT_NAMESPACE");
+    initLibNvInferPlugins(static_cast<void*>(&ctx->logger()), "");
     for (int i = 0; i < model.opset_import().size(); ++i)
     {
         std::string domain = model.opset_import(i).domain();
@@ -486,9 +510,7 @@ Status ModelImporter::importModel(
         // ONNX spec says that the default domain is either an empty string or is "ai.onnx".
         if ((domain.empty() || domain == "ai.onnx") && version < 7)
         {
-            LOG_WARNING(
-                "TensorRT supports ONNX graphs generated with at least opset 7. Models using older opsets are not "
-                "guaranteed to work.");
+            LOG_WARNING("TensorRT supports ONNX graphs generated with at least opset 7. Models using older opsets are not guaranteed to work.");
         }
         _importer_ctx.addOpset(domain, version);
     }
@@ -630,6 +652,9 @@ bool ModelImporter::parseFromFile(const char* onnxModelFile, int verbosity)
         return EXIT_FAILURE;
     }
 
+    // Keep track of the absolute path to the ONNX file.
+    _importer_ctx.setOnnxFileLocation(onnxModelFile);
+
     if (verbosity >= (int) nvinfer1::ILogger::Severity::kWARNING)
     {
         int64_t opset_version = (onnx_model.opset_import().size() ? onnx_model.opset_import(0).version() : 0);

diff --git a/ModelImporter.hpp b/ModelImporter.hpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.
+ * Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
  *
  * Permission is hereby granted, free of charge, to any person obtaining a
  * copy of this software and associated documentation files (the "Software"),

diff --git a/NvOnnxParser.cpp b/NvOnnxParser.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.
+ * Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
  *
  * Permission is hereby granted, free of charge, to any person obtaining a
  * copy of this software and associated documentation files (the "Software"),

diff --git a/NvOnnxParser.h b/NvOnnxParser.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.
+ * Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
  *
  * Permission is hereby granted, free of charge, to any person obtaining a
  * copy of this software and associated documentation files (the "Software"),
@@ -27,6 +27,12 @@
 #include <stddef.h>
 #include <vector>
 
+//!
+//! \file NvOnnxParser.h
+//!
+//! This is the API for the ONNX Parser
+//!
+
 #define NV_ONNX_PARSER_MAJOR 0
 #define NV_ONNX_PARSER_MINOR 1
 #define NV_ONNX_PARSER_PATCH 0
@@ -122,7 +128,7 @@ class IParser
      *         This method has very limited diagnostic. If parsing the serialized model
      *         fails for any reason (e.g. unsupported IR version, unsupported opset, etc.)
      *         it the user responsibility to intercept and report the error.
-     *         To obtain a better diagnostic, use the parseFromFile method below. 
+     *         To obtain a better diagnostic, use the parseFromFile method below.
      *
      * \param serialized_onnx_model Pointer to the serialized ONNX model
      * \param serialized_onnx_model_size Size of the serialized ONNX model
@@ -134,14 +140,14 @@ class IParser
                        size_t serialized_onnx_model_size)
         = 0;
 
-    /** \brief Parse an onnx model file, can be a binary protobuf or a text onnx model 
-     *         calls parse method inside.  
+    /** \brief Parse an onnx model file, can be a binary protobuf or a text onnx model
+     *         calls parse method inside.
      *
      * \param File name
      * \param Verbosity Level
-     *        
+     *
      * \return true if the model was parsed successfully
-     * 
+     *
      */
     virtual bool parseFromFile(const char* onnxModelFile, int verbosity) = 0;
 

diff --git a/OnnxAttrs.cpp b/OnnxAttrs.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.
+ * Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
  *
  * Permission is hereby granted, free of charge, to any person obtaining a
  * copy of this software and associated documentation files (the "Software"),
@@ -23,6 +23,7 @@
 #include "OnnxAttrs.hpp"
 #include "ShapedWeights.hpp"
 #include "onnx2trt_utils.hpp"
+#include <onnx/onnx_pb.h>
 
 template <>
 float OnnxAttrs::get<float>(const std::string& key) const
@@ -109,10 +110,10 @@ onnx2trt::ShapedWeights OnnxAttrs::get<onnx2trt::ShapedWeights>(const std::strin
 {
     ::ONNX_NAMESPACE::TensorProto const& onnx_weights_tensor = this->at(key)->t();
     onnx2trt::ShapedWeights weights;
-    // Return empty weights if conversion failed
-    if (!convertOnnxWeights(onnx_weights_tensor, &weights, mCtx))
+    bool success = convertOnnxWeights(onnx_weights_tensor, &weights, mCtx);
+    if (!success)
     {
-        return onnx2trt::ShapedWeights::empty(::ONNX_NAMESPACE::TensorProto_DataType_FLOAT);
+        throw std::runtime_error{"Unable to convert ONNX weights"};
     }
     return weights;
 }