自制深度学习推理框架-第五课-起飞！框架中的算子注册机制

我们的课程主页

https://github.com/zjhellofss/KuiperInfer欢迎pr和点赞

https://www.bilibili.com/vide...

为什么要有注册机制

KuiperInfer内部维护了一个注册表，用于查找Layer对应的初始化函数。这里的Layer是KuiperInfer中的算子具体执行者，例如我们在上一节课中我们讲过的ReluLayer，用于具体执行Relu方法，我们这里回顾一下Layer的定义：

class Layer {
 public:
  explicit Layer(const std::string &layer_name);

  virtual void Forwards(const std::vector<std::shared_ptr<Tensor<float>>> &inputs,
                        std::vector<std::shared_ptr<Tensor<float>>> &outputs);

  virtual ~Layer() = default;
 private:
  std::string layer_name_; 
};

Layer注册机制的实现

注册机制的原理

今天要讲的这一注册机制用到了设计模式中的工厂模式和单例模式，所以这节课也是对两大设计模式的一个合理应用和实践。KuiperInfer的注册表是一个map数据结构，维护了一组键值对，key是对应的OpType，用来查找对应的value，value是用于创建该层的对应方法(Creator)。我们可以看一下KuiperInfer中的Layer注册表实现：

 typedef std::map<OpType, Creator> CreateRegistry;

创建该层的对应方法相当于一个工厂(Creator)，Creator如下的代码所示，是一个函数指针类型，我们将存放参数的Oprator类传入到该方法中，然后该方法根据Operator内的参数返回具体的Layer.

typedef std::shared_ptr<Layer> (*Creator)(const std::shared_ptr<Operator> &op);

向注册表中注册Layer

如果我们将Layer的注册机制当成一个注册表的话，那么就会有存入的阶段也有取出的阶段，什么时候将Layer的注册机制存入到注册表呢？以如下的代码ReluLayer作为一个例子：

ReluLayer::ReluLayer(const std::shared_ptr<Operator> &op) : Layer("Relu") {
  CHECK(op->op_type_ == OpType::kOperatorRelu) << "Operator has a wrong type: " << int(op->op_type_);
  ReluOperator *relu_op = dynamic_cast<ReluOperator *>(op.get());
  CHECK(relu_op != nullptr) << "Relu operator is empty";
  this->op_ = std::make_unique<ReluOperator>(relu_op->get_thresh());
}

void ReluLayer::Forwards(const std::vector<std::shared_ptr<Tensor<float>>> &inputs,
                         std::vector<std::shared_ptr<Tensor<float>>> &outputs) {

  CHECK(this->op_ != nullptr);
  CHECK(this->op_->op_type_ == OpType::kOperatorRelu);
  const uint32_t batch_size = inputs.size();
  for (int i = 0; i < batch_size; ++i) {
    CHECK(!inputs.at(i)->empty());
    const std::shared_ptr<Tensor<float>> &input_data = inputs.at(i); 
    input_data->data().transform([&](float value) {
      float thresh = op_->get_thresh();
      if (value >= thresh) {
        return value; 
      } else {
        return 0.f;
      }
    });
    outputs.push_back(input_data);
  }
}

std::shared_ptr<Layer> ReluLayer::CreateInstance(const std::shared_ptr<Operator> &op) {
  std::shared_ptr<Layer> relu_layer = std::make_shared<ReluLayer>(op);
  return relu_layer;
}

LayerRegistererWrapper kReluLayer(OpType::kOperatorRelu, ReluLayer::CreateInstance);

LayerRegistererWrapper kReluLayer(OpType::kOperatorRelu, ReluLayer::CreateInstance), 完成了ReluLayer定义后的注册，其中key为kOperatorRelu, value 为ReluLayer::CreateInstance. CreateInstance是一个具体的工厂方法，用来在之后对ReluLayer进行初始化，我们接下来看看这里注册机制的具体实现：

class LayerRegistererWrapper {
 public:
  LayerRegistererWrapper(OpType op_type, const LayerRegisterer::Creator &creator) {
    LayerRegisterer::RegisterCreator(op_type, creator);
  }
};

LayerRegistererWrapper是一个类

我们在调用kReluLayer(OpType::kOperatorRelu, ReluLayer::CreateInstance)的时候，LayerRegistererWrapper的构造方法再调用了RegisterCreator。所以到目前为止，调用链是这样的：

ReluLayer定义完成--->LayerRegistererWrapper ---> RegisterCreator

接下来看看RegisterCreator这个注册方法的实现：

void LayerRegisterer::RegisterCreator(OpType op_type, const Creator &creator) {
  CHECK(creator != nullptr) << "Layer creator is empty";
  CreateRegistry &registry = Registry();
  CHECK_EQ(registry.count(op_type), 0) << "Layer type: " << int(op_type) << " has already registered!";
  registry.insert({op_type, creator});
}

再强调一遍，方法中的op_type是算子的类型，作为Layer注册表的key使用，creator是创建具体层的工厂方法，作为Layer注册表的value. 目前为止，调用链是这样的：

ReluLayer定义完成 --->LayerRegistererWrapper ---> RegisterCreator --->Registry返回注册表 --->存入实现方法

单例注册表的实现

可以看到CreateRegistry &registry =Registry()这里返回注册表的实例，此处的Layer注册表是全局唯一的，全局唯一的实现方法是单例设计模式的应用，我们看一下下方的具体实现：

LayerRegisterer::CreateRegistry &LayerRegisterer::Registry() {
  static CreateRegistry *kRegistry = new CreateRegistry();
  CHECK(kRegistry != nullptr) << "Global layer register init failed!";
  return *kRegistry;
}

此处的static CreateRegistry *kRegistry =newCreateRegistry() , 使得Layer注册表在全局有且只有一个，无论我们调用了多少次Registry(),该方法都会返回同一个Layer注册表。

向注册表中取出Layer

在完成Layer注册之后，我们就可以根据OpType来取出用于实例化一个Layer的工厂函数，比如上面我们完成了ReluLayer的注册后，系统中的Layer注册表中是这样的:

{kOperatorRelu:ReluLayer::CreateInstance}

我们在需要的时候时候根据kOperator来取出ReluLayer::CreateInstance, 我们再以Relu中的工厂函数为例子看看一个具体工厂函数的实现：

std::shared_ptr<Layer> ReluLayer::CreateInstance(const std::shared_ptr<Operator> &op) {
  std::shared_ptr<Layer> relu_layer = std::make_shared<ReluLayer>(op);
  return relu_layer;
}

可以看到这里的工厂函数比较简单，直接传入ReluOperator完成对ReluLayer的初始化。

一个例子

TEST(test_layer, forward_relu2) {
  using namespace kuiper_infer;
  float thresh = 0.f;
  std::shared_ptr<Operator> relu_op = std::make_shared<ReluOperator>(thresh);
  std::shared_ptr<Layer> relu_layer = LayerRegisterer::CreateLayer(relu_op);

  std::shared_ptr<Tensor<float>> input = std::make_shared<Tensor<float>>(1, 1, 3);
  input->index(0) = -1.f;
  input->index(1) = -2.f;
  input->index(2) = 3.f;
  std::vector<std::shared_ptr<Tensor<float>>> inputs;
  std::vector<std::shared_ptr<Tensor<float>>> outputs;
  inputs.push_back(input);
  relu_layer->Forwards(inputs, outputs);
  ASSERT_EQ(outputs.size(), 1);
  for (int i = 0; i < outputs.size(); ++i) {
    ASSERT_EQ(outputs.at(i)->index(0), 0.f);
    ASSERT_EQ(outputs.at(i)->index(1), 0.f);
    ASSERT_EQ(outputs.at(i)->index(2), 3.f);
  }
}

我们可以看到std::shared_ptr<Operator> relu_op = std::make_shared<ReluOperator>(thresh), 初始化了一个ReluOperator, 其中的参数为thresh=0.f.

因为我们已经在ReluLayer的实现中完成了注册，{kOperatorRelu:ReluLayer::CreateInstance}, 所以现在可以使用LayerRegisterer::CreateLayer(relu_op)得到我们ReluLayer中的实例化工厂方法，我们再来看看CreateLayer的实现：

std::shared_ptr<Layer> LayerRegisterer::CreateLayer(const std::shared_ptr<Operator> &op) {
  CreateRegistry &registry = Registry();
  const OpType op_type = op->op_type_;

  LOG_IF(FATAL, registry.count(op_type) <= 0) << "Can not find the layer type: " << int(op_type);
  const auto &creator = registry.find(op_type)->second;

  LOG_IF(FATAL, !creator) << "Layer creator is empty!";
  std::shared_ptr<Layer> layer = creator(op);
  LOG_IF(FATAL, !layer) << "Layer init failed!";
  return layer;
}

可以看到传入的参数为op, 我们首先取得op中的op\_type, 此处的op\_type为kOperatorRelu, 根据registry.find(op_type), 就得到了层的初始化方法creator, 随后使用传入的op去初始化layer并返回实例。值得注意的是此处也调用了CreateRegistry &registry =Registry()返回了我们所说的全局有且唯一的Layer注册表. 此处的creator(op)就相当于调用了ReluLayer::CreateInstance.