NvOnnxParser.h

前往此文件的文档。

/*
 * 版权所有 (c) 1993-2024, NVIDIA CORPORATION。保留所有权利。
 *
 * 特此授权，免费授予任何获得
 * 本软件和相关文档文件（以下简称“软件”）副本的人员
处理本软件不受限制，包括但不限于
使用、复制、修改、合并、发布、分发、再许可的权利，
和/或出售本软件副本的权利，并允许向其
提供本软件的人员这样做，但须遵守以下条件：
 *
* 上述版权声明和本许可声明应包含在
* 本软件的所有副本或主要部分中。
 *
* 本软件按“现状”提供，不作任何明示或暗示的担保，包括但不限于
* 对适销性、特定用途的适用性和非侵权性的保证。
* 在任何情况下，作者或版权持有人均不对任何索赔、损害或其他
* 责任负责，无论是在合同、侵权行为或其他方面引起的诉讼中，
* 因本软件或本软件的使用或其他
* 与本软件相关的交易而产生。
 * DEALINGS IN THE SOFTWARE.
 */
 
#ifndef NV_ONNX_PARSER_H
#define NV_ONNX_PARSER_H
 
#include "NvInfer.h"
#include <stddef.h>
#include <string>
#include <vector>
 
 
#define NV_ONNX_PARSER_MAJOR 0
#define NV_ONNX_PARSER_MINOR 1
#define NV_ONNX_PARSER_PATCH 0
 
static constexpr int32_t NV_ONNX_PARSER_VERSION
 = ((NV_ONNX_PARSER_MAJOR * 10000) + (NV_ONNX_PARSER_MINOR * 100) + NV_ONNX_PARSER_PATCH);
 
typedef std::pair<std::vector<size_t>, bool> SubGraph_t;
 
typedef std::vector<SubGraph_t> SubGraphCollection_t;
 
namespace nvonnxparser
{
 
template <typename T>
constexpr inline int32_t EnumMax() noexcept;
 
enum class ErrorCode : int
{
kSUCCESS = 0, // 成功
kINTERNAL_ERROR = 1, // 内部错误
kMEM_ALLOC_FAILED = 2, // 内存分配失败
kMODEL_DESERIALIZE_FAILED = 3, // 模型反序列化失败
kINVALID_VALUE = 4, // 无效值
kINVALID_GRAPH = 5, // 无效图
kINVALID_NODE = 6, // 无效节点
kUNSUPPORTED_GRAPH = 7, // 不支持的图
kUNSUPPORTED_NODE = 8, // 不支持的节点
kUNSUPPORTED_NODE_ATTR = 9, // 不支持的节点属性
kUNSUPPORTED_NODE_INPUT = 10, // 不支持的节点输入
kUNSUPPORTED_NODE_DATATYPE = 11, // 不支持的节点数据类型
kUNSUPPORTED_NODE_DYNAMIC = 12, // 不支持的节点动态性
kUNSUPPORTED_NODE_SHAPE = 13, // 不支持的节点形状
kREFIT_FAILED = 14 // 重拟合失败
};
 
template <>
constexpr inline int32_t EnumMax<ErrorCode>() noexcept
{
 return 14;
}
 
using OnnxParserFlags = uint32_t;
 
enum class OnnxParserFlag : int32_t
{
kNATIVE_INSTANCENORM = 0 // 本地 InstanceNorm
};
 
template <>
constexpr inline int32_t EnumMax<OnnxParserFlag>() noexcept
{
 return 1;
}
 
class IParserError
{
public
 virtual ErrorCode code() const = 0;
 virtual char const* desc() const = 0;
 virtual char const* file() const = 0;
 virtual int line() const = 0;
 virtual char const* func() const = 0;
 virtual int node() const = 0;
 virtual char const* nodeName() const = 0;
 virtual char const* nodeOperator() const = 0;
 virtual char const* const* localFunctionStack() const = 0;
// 堆栈大小为 0。
 virtual int32_t localFunctionStackSize() const = 0;
 
protected
 virtual ~IParserError() {}
};
 
class IParser
{
public
 virtual bool parse(
 void const* serialized_onnx_model, size_t serialized_onnx_model_size, const char* model_path = nullptr) noexcept
        = 0;
 
 virtual bool parseFromFile(const char* onnxModelFile, int verbosity) noexcept = 0;
 
 TRT_DEPRECATED virtual bool supportsModel(void const* serialized_onnx_model, size_t serialized_onnx_model_size,
 SubGraphCollection_t& sub_graph_collection, const char* model_path = nullptr) noexcept = 0;
 
 virtual bool parseWithWeightDescriptors(
 void const* serialized_onnx_model, size_t serialized_onnx_model_size) noexcept
        = 0;
 
 virtual bool supportsOperator(const char* op_name) const noexcept = 0;
 
 virtual int getNbErrors() const noexcept = 0;
 
 virtual IParserError const* getError(int index) const noexcept = 0;
 
 virtual void clearErrors() noexcept = 0;
 
 virtual ~IParser() noexcept = default;
 
 virtual char const* const* getUsedVCPluginLibraries(int64_t& nbPluginLibs) const noexcept = 0;
 
 virtual void setFlags(OnnxParserFlags onnxParserFlags) noexcept = 0;
 
 virtual OnnxParserFlags getFlags() const noexcept = 0;
 
 virtual void clearFlag(OnnxParserFlag onnxParserFlag) noexcept = 0;
 
 virtual void setFlag(OnnxParserFlag onnxParserFlag) noexcept = 0;
 
 virtual bool getFlag(OnnxParserFlag onnxParserFlag) const noexcept = 0;
 
 virtual nvinfer1::ITensor const* getLayerOutputTensor(char const* name, int64_t i) noexcept = 0;
 
 virtual bool supportsModelV2(
 void const* serializedOnnxModel, size_t serializedOnnxModelSize, char const* modelPath = nullptr) noexcept = 0;
 
 virtual int64_t getNbSubgraphs() noexcept = 0;
 
 virtual bool isSubgraphSupported(int64_t const index) noexcept = 0;
 
 virtual int64_t* getSubgraphNodes(int64_t const index, int64_t& subgraphLength) noexcept = 0;
};
 
class IParserRefitter
{
public
 virtual bool refitFromBytes(
 void const* serializedOnnxModel, size_t serializedOnnxModelSize, char const* modelPath = nullptr) noexcept
        = 0;
 
 virtual bool refitFromFile(char const* onnxModelFile) noexcept = 0;
 
 virtual int32_t getNbErrors() const noexcept = 0;
 
 virtual IParserError const* getError(int32_t index) const noexcept = 0;
 
 virtual void clearErrors() = 0;
 
 virtual ~IParserRefitter() noexcept = default;
};
 
} // namespace nvonnxparser
 
extern "C" TENSORRTAPI void* createNvOnnxParser_INTERNAL(void* network, void* logger, int version) noexcept;
extern "C" TENSORRTAPI void* createNvOnnxParserRefitter_INTERNAL(
 void* refitter, void* logger, int32_t version) noexcept;
extern "C" TENSORRTAPI int getNvOnnxParserVersion() noexcept;
 
namespace nvonnxparser
{
 
namespace
{
 
inline IParser* createParser(nvinfer1::INetworkDefinition& network, nvinfer1::ILogger& logger) noexcept
{
 try
    {
 return static_cast<IParser*>(createNvOnnxParser_INTERNAL(&network, &logger, NV_ONNX_PARSER_VERSION));
    }
 catch (std::exception& e)
    {
 logger.log(nvinfer1::ILogger::Severity::kINTERNAL_ERROR, e.what());
    }
 
 return nullptr;
}
 
inline IParserRefitter* createParserRefitter(nvinfer1::IRefitter& refitter, nvinfer1::ILogger& logger) noexcept
{
 try
    {
 return static_cast<IParserRefitter*>(
 createNvOnnxParserRefitter_INTERNAL(&refitter, &logger, NV_ONNX_PARSER_VERSION));
    }
 catch (std::exception& e)
    {
 logger.log(nvinfer1::ILogger::Severity::kINTERNAL_ERROR, e.what());
    }
 
 return nullptr;
}
 
} // namespace
 
} // namespace nvonnxparser
 
#endif // NV_ONNX_PARSER_H