[源码阅读] TensorRT - ONNX parser

相关内容：
[源码阅读] 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;
}

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