[源码阅读] TensorRT - ONNX parser

本文详细介绍了ONNXparser在TensorRT中的实现,包括如何通过映射函数将ONNX模型转换为TensorRT网络定义。重点讨论了ShapedWeights类,IImporterContext接口,以及NodeImporter函数模板,用于将ONNX算子转化为TensorRT的激活层。此外,还提到了OnnxAttrs类用于处理ONNX模型的属性,并展示了如何处理不同数据类型的属性值。

摘要生成于 C知道 ,由 DeepSeek-R1 满血版支持, 前往体验 >

文章目录


相关内容:
[源码阅读] TensorRT —— Caffe Parser、Plugin

ONNX parser

TensorRT 的 onnx parser 是一个独立的开源项目。 onnx-tensorrt:https://github.com/onnx/onnx-tensorrt

这里映射关系还是首先建立起来一个个映射函数,然后通过 op_type 来找到对应的映射函数。

typedef ValueOrStatus<std::vector<TensorOrWeights>> NodeImportResult;

// 映射函数别名
typedef std::function<NodeImportResult(IImporterContext *ctx,
                                       ::ONNX_NAMESPACE::NodeProto const &node,
                                       std::vector<TensorOrWeights> &inputs)>
    NodeImporter;

可以看到映射函数的主体是 返回类型为 NodeImportResult, 输入参数为 IImporterContext 类型:


class ShapedWeights {
public:
  using DataType = int32_t;
  DataType type;
  void* values;
  nvinfer1::Dims shape;
  static ShapedWeights empty(DataType type);
  ShapedWeights();
  explicit ShapedWeights(DataType type, void* values, nvinfer1::Dims shape_);
  size_t count() const;
  size_t size_bytes() const;
  operator bool() const;
  operator nvinfer1::Weights() const;
};

struct IImporterContext {
  virtual nvinfer1::INetworkDefinition* network() = 0;
  virtual ShapedWeights createTempWeights(ShapedWeights::DataType type,
                                          nvinfer1::Dims shape) = 0;
  virtual int64_t getOpsetVersion(const char* domain="") const = 0;
protected:
  virtual ~IImporterContext() {}
};

其中包含了一个 INetworkDefinition 对象来使用 API 构建 network,一个创建 ShapedWeights 的成员函数,还有一个获取 OpsetVersion 的接口。

注意这里为了方便直接将 SharedWeights 类直接送入 INetworkDefinition 的接口, 这里使用了 C++隐式类型转换运算符operator type()用法,可以在需要的时候将 SharedWeights 类隐式转换成 TensorRT 需要的 Weights 类型。相关映射实现为:

ShapedWeights::operator nvinfer1::Weights() const {
  nvinfer1::Weights w{};
  // If INT64 weights, check if all the values can be cast down to INT32.
  if (this->type == ::ONNX_NAMESPACE::TensorProto::INT64) {
    std::vector<int32_t> int32_weights;
    int32_weights.resize(this->count());

    if (!onnx2trt::convertINT64(this->values, this->count(), int32_weights)) {
      // Return empty w on failure
      return w;
    } else {
      void * int32_weights_ptr = static_cast<void *>(int32_weights.data());
      std::memcpy(this->values, int32_weights_ptr, int32_weights.size() * sizeof(int32_t));
      w.values = this->values;
    }
  } else {
    w.values = this->values;
  }
  bool supported_type = convert_dtype(this->type, &w.type);
  (void)supported_type;
  assert(supported_type);
  w.count = this->count();
  return w;
}

NodeProto 类型,以及一个代表数据的 vector, 其中:


class TensorOrWeights {
  union {
    nvinfer1::ITensor* _tensor;
    ShapedWeights      _weights;
  };
  enum { NODE_TENSOR, NODE_WEIGHTS } _variant;
public:
  inline TensorOrWeights() : _tensor(nullptr), _variant(NODE_TENSOR) {}
  inline TensorOrWeights(nvinfer1::ITensor* tensor)
    : _tensor(tensor), _variant(NODE_TENSOR) {}
  inline TensorOrWeights(ShapedWeights const& weights)
    : _weights(weights), _variant(NODE_WEIGHTS) {}
  inline bool is_tensor()  const { return _variant == NODE_TENSOR; }
  inline bool is_weights() const { return _variant == NODE_WEIGHTS; }
  inline nvinfer1::ITensor& tensor() {
    assert(this->is_tensor());
    return *_tensor;
  }
  inline nvinfer1::ITensor const& tensor() const {
    assert(this->is_tensor());
    return *_tensor;
  }
  inline ShapedWeights& weights() {
    assert(this->is_weights());
    return _weights;
  }
  inline ShapedWeights const& weights() const {
    assert(this->is_weights());
    return _weights;
  }
  inline nvinfer1::Dims shape() const {
    return this->is_tensor() ? _tensor->getDimensions() : _weights.shape;
  }
  inline operator bool() const {
    return this->is_tensor() ? (bool)_tensor : (bool)_weights;
  }
  nvinfer1::ITensor* reset_tensor(nvinfer1::ITensor* tensor) {
    assert(this->is_tensor());
    return _tensor = tensor;
  }
};

上面看到是使用联合体来表示 TensorOrWeights 这个类的。 基于上面的一些基础设施,onnx-tensorrt 写了一系列的算子加载函数。

全在: onnx-tensorrt/builtin_op_importers.cpp 中。 通过宏,都会将它们注册到最终的一个 map 中去。

string_map<NodeImporter>& getBuiltinOpImporterMap()
{
    static string_map<NodeImporter> builtin_op_importers;
    return builtin_op_importers;
}

bool registerBuiltinOpImporter(std::string op, NodeImporter const& importer)
{
    bool inserted = getBuiltinOpImporterMap().insert({op, importer}).second;
    assert(inserted);
    return inserted;
}

#define DEFINE_BUILTIN_OP_IMPORTER(op)                                                                                 \
    NodeImportResult import##op(                                                                                       \
        IImporterContext* ctx, ::ONNX_NAMESPACE::NodeProto const& node, std::vector<TensorOrWeights>& inputs);         \
    static const bool op##_registered_builtin_op = registerBuiltinOpImporter(#op, import##op);                         \
    IGNORE_UNUSED_GLOBAL(op##_registered_builtin_op);                                                                  \
    NodeImportResult import##op(                                                                                       \
        IImporterContext* ctx, ::ONNX_NAMESPACE::NodeProto const& node, std::vector<TensorOrWeights>& inputs)

采用上述方式,这里以 LeakyRelu 为例:

DEFINE_BUILTIN_OP_IMPORTER(LeakyRelu)
{
    OnnxAttrs attrs(node);
    float alpha = attrs.get<float>("alpha", 0.01f); // 处理相关的属性
    return activationHelper(ctx, node, inputs, nvinfer1::ActivationType::kLEAKY_RELU, &alpha);
}

NodeImportResult activationHelper(IImporterContext* ctx, const ::ONNX_NAMESPACE::NodeProto& node,
    std::vector<TensorOrWeights>& inputs, nvinfer1::ActivationType op, float* alpha, float* beta)
{
    nvinfer1::ITensor& input = convertToTensor(inputs.at(0), ctx);
    nvinfer1::IActivationLayer* layer = ctx->network()->addActivation(input, op);
    if (alpha)
    {
        layer->setAlpha(*alpha);
    }
    if (beta)
    {
        layer->setBeta(*beta);
    }

    return {{layer->getOutput(0)}};
}

本身 ONNX 算子可能会有 opset 的问题,不同 opset 版本的同一个算子可能会有不同的属性或者不同个数的输入,因为 TensorRT 这个逐个算子处理映射,在内部逻辑里面就把这些差异直接内部处理了,最后转化成 TensorRT 中比较统一的形式。

class OnnxAttrs {
  template<typename T>
  using string_map = std::unordered_map<std::string, T>;
  typedef string_map<::ONNX_NAMESPACE::AttributeProto const*> AttrMap;
  AttrMap _attrs;
public:
  explicit OnnxAttrs(::ONNX_NAMESPACE::NodeProto const& onnx_node) {
    for( auto const& attr : onnx_node.attribute() ) {
      _attrs.insert({attr.name(), &attr});
    }
  }
  bool count(std::string key) const { return _attrs.count(key); }
  ::ONNX_NAMESPACE::AttributeProto const* at(std::string key) const {
    if( !_attrs.count(key) ) {
      throw std::out_of_range("Attribute not found: " + key);
    }
    return _attrs.at(key);
  }
  template<typename T> T get(const std::string& key) const;
  template<typename T> T get(const std::string& key, T const& default_value) const {
    return _attrs.count(key) ? this->get<T>(key) : default_value;
  }
};

属性的处理主要是OnnxAttrs 类,通过 NodeProto 来初始化一个 OnnxAttrs 类,初始化的过程中就将所有的 AttributeProto, 按照 {name : AttributeProto} 的形式全部 push 到一个 Map 里,在获取的时候会首先看在 map 里是否可以找到,找不到就直接报错了。

上面使用了模板成员函数,然后实例化了一下相关会用到的数据类型的属性获取,方便后面使用。

// 一些基础类型的属性值获取
template<> bool OnnxAttrs::get<bool>(const std::string& key) const {
  int value = this->at(key)->i();
  assert(value == bool(value));
  return bool(value);
}

template<> std::string OnnxAttrs::get<std::string>(const std::string& key) const {
  return this->at(key)->s();
}

template<> std::vector<int> OnnxAttrs::get<std::vector<int>>(const std::string& key) const {
  auto attr = this->at(key)->ints();
  return std::vector<int>(attr.begin(), attr.end());
}

// 将属性直接映射到 TensorRT 的数据结构上
template<> nvinfer1::Dims OnnxAttrs::get<nvinfer1::Dims>(const std::string& key) const {
  auto values = this->get<std::vector<int>>(key);
  nvinfer1::Dims dims;
  dims.nbDims = values.size();
  std::copy(values.begin(), values.end(), dims.d);
  // Note: No dimension type information is included
  return dims;
}

template<> onnx2trt::ShapedWeights OnnxAttrs::get<onnx2trt::ShapedWeights>(const std::string& key) const {
  ::ONNX_NAMESPACE::TensorProto const& onnx_weights_tensor = this->at(key)->t();
  onnx2trt::ShapedWeights weights;
  convert_onnx_weights(onnx_weights_tensor, &weights);
  return weights;
}

这样使用模板成员函数,可以一定程度上消除重复代码,因为几乎每一个算子都需要处理不同数据类型的属性值。

### NVIDIA TensorRT 库使用指南 #### 安装与配置 NVIDIA® TensorRT™ 是专为在 NVIDIA GPU 上实现高性能深度学习推理而设计的 SDK。为了顺利安装并配置该工具包,需遵循官方文档中的指导说明[^1]。 对于大多数用户而言,推荐采用预编译二进制文件的方式来进行安装。这可以通过下载适用于目标平台(Linux 或 Windows)的 `.deb` 或者 `.rpm` 文件来完成。另一种方法是从源码构建,这种方式适合那些希望自定义设置或需要特定功能支持的人群。无论哪种方式,在开始之前都应确认已正确设置了 CUDA 和 cuDNN 环境。 一旦完成了基本环境搭建之后,则可通过 Python API 来简化集成过程: ```bash pip install nvidia-tensorrt ``` 上述命令会自动处理依赖关系并将必要的库添加到环境中[^2]。 #### 性能优化策略 针对模型部署阶段遇到的各种挑战,TensorRT 提供了一系列性能调优手段以确保最佳运行效率。其中包括但不限于以下几个方面: - **INT8量化**:通过减少权重精度至8位整数表示形式,可以在几乎不影响准确性的情况下显著提升吞吐量; - **FP16半精度浮点运算**:利用GPU硬件特性加速计算速度的同时降低功耗; - **多流并发执行模式**:允许单个设备上同时处理多个输入批次的数据请求,从而提高整体利用率; - **动态形状支持**:使得网络结构能够适应不同尺寸的输入图像而不必重新训练整个模型。 下面给出一段简单的Python脚本作为示例展示如何加载预先训练好的ONNX格式模型并通过TensorRT进行推断操作: ```python import tensorrt as trt from polygraphy.backend.trt import CreateConfig, EngineFromNetwork, TrtRunner def build_engine(onnx_file_path): network_creation_flag = 0 with open(onnx_file_path,'rb') as model: builder = trt.Builder(trt.Logger()) explicit_batch = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH) network = builder.create_network(explicit_batch) parser = trt.OnnxParser(network,trt.Logger()) if not parser.parse(model.read()): print('Failed to parse the ONNX') for error in range(parser.num_errors): print(parser.get_error(error)) return None config = builder.create_builder_config() profile = builder.create_optimization_profile() engine = builder.build_serialized_network(network,config) return engine engine = build_engine("model.onnx") with TrtRunner(EngineFromNetwork(engine)) as runner: outputs = runner.infer(feed_dict={"input": input_data}) print(outputs) ```
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值