注册机调研：从Caffe开始

一、简介

现在开源的深度学习框架都会用到注册机，我准备做一个系列，来详细看一下各个开源框架注册机的实现，先从Caffe开始。Caffe代码现在来看虽然比较古老，但是作为训练框架的先驱，依然有值得学习的地方，比如模块化、高性能（单机单卡）、易扩展等，本文着重来看里面的注册机制。

caffe里主要有两个注册宏：

• REGISTER_LAYER_CLASS

• REGISTER_SOLVER_CLASS

下面通过实际代码来详细解读这两个注册宏（注：为方便理解，相关代码都做了简化和修改）

二、REGISTER_LAYER_CLASS

这个注册宏用于注册一个函数，这个函数用于创建某种layer的对象，这样在创建layer的时候，代码会比较简洁容易维护，不易出错。

而如果使用类似下面这种if或者switch的做法来创建layer，当layer种类少的时候还没问题，当layer种类越来越多的时候，不光性能不好，也很容易写错，现在Caffe的官方仓库里有大概五十种layer，试想如果都在下面列出来，很容易在某个if语句上就会写错。

if (layer_type == "conv1d") { 
  return MakeConv1dLayer(param); 
} else if (layer_type == "conv2d") {
  return MakeConv2dLayer(param);
} else if (layer_type == "conv3d") {
  return MakeConv3dLayer(param);
} else if(...) {
  ...
}

如果使用注册的方式，就不会有上面所说的问题，一个Layer实现在一个单独的文件中，在文件的末尾来做一次注册，不会让人眼花缭乱，大大减少笔误的可能性：

// 以concat_layer.cpp为例，前面是实现，最后加一句注册即可
// concat实现部分
...
void ConcatLayer::Forward_cpu(...) {...}
void ConcatLayer::Backward_cpu(...) {...}
// 注册
REGISTER_LAYER_CLASS(Concat);

另外由于注册机中的数据结构为map，在创建某个给定type的layer时，时间复杂度为O(1)，而前面伪代码的时间复杂度为O(n)。

下面来详细看下具体的实现，这个宏的定义在layer_factory.hpp文件中，先看具体定义：

#define REGISTER_LAYER_CREATOR(type, creator)                                  \
  static LayerRegisterer<float> g_creator_f_##type(#type, creator<float>);     \
  static LayerRegisterer<double> g_creator_d_##type(#type, creator<double>)    \

#define REGISTER_LAYER_CLASS(type)                                             \
  template <typename Dtype>                                                    \
  shared_ptr<Layer<Dtype>> Creator_##type##Layer(const LayerParameter& param){ \
    return shared_ptr<Layer<Dtype>>(new type##Layer<Dtype>(param));            \
  }                                                                            \
  REGISTER_LAYER_CREATOR(type, Creator_##type##Layer)

其中LayerRegisterer是个模板类，构造函数参数是layer type和一个函数指针，在调用宏注册Layer的时候，会定义两个LayerRegisterer类型的全局变量，这样LayerRegisterer的构造函数就会被自动执行：

template <typename Dtype>
class LayerRegisterer {
 public:
  LayerRegisterer(const string& type,
      shared_ptr<Layer<Dtype>> (*creator)(const LayerParameter&)) {
    LayerRegistry<Dtype>::AddCreator(type, creator);
  }
};

从上面代码可见LayerRegistry的AddCreator被调用执行，LayerRegistry中由三个主要的静态函数组成：

• 一个Registry函数，用于返回一个全局的存储数据结构

• 一个AddCreator函数，用于注册一个函数

• 一个CreateLayer函数，用于调用注册的函数

精简代码如下：

// include/caffe/layer_factory.hpp
template <typename Dtype>
class LayerRegistry {
 public:
  typedef shared_ptr<Layer<Dtype>> 
      (*Creator)(const LayerParameter&);
  typedef std::map<string, Creator> CreatorRegistry;

  static CreatorRegistry& Registry() {
    static CreatorRegistry* g_registry_ = new CreatorRegistry();
    return *g_registry_;
  }
  static void AddCreator(const string& type, Creator creator) {
    CreatorRegistry& registry = Registry();
    registry[type] = creator;
  }
  static shared_ptr<Layer<Dtype>> CreateLayer(
        const LayerParameter& param) {
    CreatorRegistry& registry = Registry();
    return registry[param.type()](param);
  }
};

在Caffe中，在运行一个实际的网络之前，需要把所有的Layer都创建出来，在Net类的Init函数中可以看到创建Layer的代码，只有一句话：

// src/caffe/net.cpp Init函数中的创建Layer代码
layers_.push_back(LayerRegistry<Dtype>::CreateLayer(layer_param));
// include/caffe/net.hpp中的layers_定义
vector<shared_ptr<Layer<Dtype>>> layers_;

三、REGISTER_SOLVER_CLASS

这个注册宏是用来注册solver的，solver的作用是定义参数的更新方式，常见的有sgd、adam等方法，其中sgd算法是我们常用的，这里只关注注册部分

在Caffe中，solver的注册实现和layer是完全一样的，只有名字上的区别，具体可以查看include/caffe/solver_factory.hpp这个文件，这里不再赘述了