caffe注册机制浅析——宏与类的使用

一、楔子

15年那会儿，我刚入坑深度学习。那时候几大主流框架caffe，theano，和torch7（不是pytorch）分别代表了C++、python、Lua几大语言。其中尤以caffe最为风靡也最受欢迎。因为我本科是电子工程，不是计算机科班出身，在甫一阅读caffe的代码，尤其是看到众多我不熟悉的概念，如常引用，模板类，protobuf，glog，gtest等等时，自然是一脸懵逼，直接劝退。当然后面的tensorflow，pytorch等框架横空出世，编程便不再成为科研问题，这是后话。

工作后，由于接手的代码全为C++，加之我个人是个很执着的人，我竟又想到了当年的“白月光”，不得不说，没把caffe玩明白一直是我心里一大遗憾。所以也想借着休息时间，再乘时光机回到6年前，读一读当时没有读懂的代码，也了了一桩心愿罢。

二、简介

1. 源码位置

本文主要关注的是caffe一个小的方面，即注册机制（Registration）。在caffe中，涉及到注册的部分主要有：

• /include/caffe/solver_factory.hpp
• /include/caffe/layer_factory.hpp
• /tools/caffe.cpp

2. 何为注册

简单谈谈注册。其实很简单，以python语言为例，我的理解就是有一个全局的字典dict，每写一个新类，就把（类的名字，类）给添加到dict里面去。这样后面给出一个字符串string，我就能得到字符串对应的类实例，如：

REGISTERER = dict()

# Define
class A(object):
    def run(self):
        print(f"Run class: {self.__class__.__name__}")
# Register
REGISTERER[A.__name__] = A


# Define
class B(object):
    def run(self):
        print(f"Run class: {self.__class__.__name__}")
# Register
REGISTERER[B.__name__] = B


if __name__ == "__main__":
    instance = REGISTERER["B"]()
    instance.run()

从上面这个例子可以看出，如果我把和字典相关的几行撤掉，换做 if - elif - else，可能也没啥问题。是的，注册就是完成这么一回事情。当然在C++中代码有不同的写法，下面我们来看看C++是如何实现的。

三、源码分析

主要分析layer_factory.hpp中的写法，会顺便带一下solver_factory.hpp和caffe.cpp中的东西，都是一样的原理。有条件的同学打开代码地址，边看源码边看分析。

首先，layer_factory.hpp是一个头文件，所有的声明都包含在caffe这个namespace下。

其次，头文件中包含了两个模板类LayerRegistry和LayerRegisterer，以及两个宏REGISTER_LAYER_CREATOR和REGISTER_LAYER_CLASS。这是代码的基本结构。下面分别介绍这4个东东。

1. LayerRegistry

先说第一个模板类LayerRegistry，这个类是主要干活的类。

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_;
  }

  // Adds a creator.
  static void AddCreator(const string& type, Creator creator) {
    CreatorRegistry& registry = Registry();
    CHECK_EQ(registry.count(type), 0)
        << "Layer type " << type << " already registered.";
    registry[type] = creator;
  }

  // Get a layer using a LayerParameter.
  static shared_ptr<Layer<Dtype> > CreateLayer(const LayerParameter& param) {
    if (Caffe::root_solver()) {
      LOG(INFO) << "Creating layer " << param.name();
    }
    const string& type = param.type();
    CreatorRegistry& registry = Registry();
    CHECK_EQ(registry.count(type), 1) << "Unknown layer type: " << type
        << " (known types: " << LayerTypeListString() << ")";
    return registry[type](param);
  }

  static vector<string> LayerTypeList() {
    CreatorRegistry& registry = Registry();
    vector<string> layer_types;
    for (typename CreatorRegistry::iterator iter = registry.begin();
         iter != registry.end(); ++iter) {
      layer_types.push_back(iter->first);
    }
    return layer_types;
  }

 private:
  // Layer registry should never be instantiated - everything is done with its
  // static variables.
  LayerRegistry() {}

  static string LayerTypeListString() {
    vector<string> layer_types = LayerTypeList();
    string layer_types_str;
    for (vector<string>::iterator iter = layer_types.begin();
         iter != layer_types.end(); ++iter) {
      if (iter != layer_types.begin()) {
        layer_types_str += ", ";
      }
      layer_types_str += *iter;
    }
    return layer_types_str;
  }
};

看着有点长，我们一点点地来。

typedef shared_ptr<Layer<Dtype> > (*Creator)(const LayerParameter&);
typedef std::map<string, Creator> CreatorRegistry;

• 首先使用typedef定义了两个类型。Creator是一个函数指针，它代表一类函数，接收LayerParameter类对象，返回一个Layer<Dtype>的smart pointer；其次，定义了一种字典类型CreatorRegistry，它的key是字符串，value是一个函数指针。怎么样？是不是和上面python代码里面的字典对上号了？

static CreatorRegistry& Registry() {
    static CreatorRegistry* g_registry_ = new CreatorRegistry();
    return *g_registry_;
}

• LayerRegistry中的所有成员方法均为静态方法，也就是说这些函数并不与类实例对象所绑定，而是直接绑在了类上，此其一。注意到Registry函数内部是定义了一个静态的CreatorRegistry*的变量。静态变量有个特点，即它的生存期是持续到程序结束的，这是因为它并不存在与堆区或者栈区，在可重定向文件，也就是.o文件中，会单独为静态变量以及全局变量建立符号表。

    回过来，这个函数的意思是，如果第一次执行那么new一个对象，而如果是第二次即以后执行，那么就使用之前new出来的对象，很明显，这就是一种典型的单例模型的写法。还要竹弈一个细节，这个函数的返回值是一个引用，也就是我们就是想拿到函数内作用域的那个对象，而非进行拷贝，这点也很关键。

// Adds a creator.
  static void AddCreator(const string& type, Creator creator) {
    CreatorRegistry& registry = Registry();
    CHECK_EQ(registry.count(type), 0)
        << "Layer type " << type << " already registered.";
    registry[type] = creator;
  }

• AddCreator仍然是一个静态的成员方法，首先调用上面的Registry函数，拿到字典，然后就是把字符串type和函数指针creator保存到字典里面。

// Get a layer using a LayerParameter.
  static shared_ptr<Layer<Dtype> > CreateLayer(const LayerParameter& param) {
    if (Caffe::root_solver()) {
      LOG(INFO) << "Creating layer " << param.name();
    }
    const string& type = param.type();
    CreatorRegistry& registry = Registry();
    CHECK_EQ(registry.count(type), 1) << "Unknown layer type: " << type
        << " (known types: " << LayerTypeListString() << ")";
    return registry[type](param);
  }

• CreateLayer的函数实现的目的是，给定LayerParameter类对象，拿到这个parameter里面的type，去之前的字典里面查询有无，如果有，那么就把函数指针registry[type]给取出来，再把param传给这个函数，返回最终的Layer对象的smart pointer。

[小结]

基本上就3个东西，新建字典，插入字典，查询字典。

2. LayerRegisterer

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

非常小的一个类，只有一个构造函数，意味着只在实例化时才起作用。其实实际上也完全可以不用类的方式，而是简单地写个函数达到同样效果。

可以看到它真正的作用是调用了AddCreator，也就是上面所说的把字符串和函数指针给加入到字典里面去。如果是第一次执行，那么LayerRegistry会自动new一个字典出来，后面再执行，都是在那个老字典里面add了。

3. 两个宏

#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)

两个宏的关系，下面的REGISTER_LAYER_CLASS调用了上面的REGISTER_LAYER_CREATOR，并且是对外使用的宏。宏的两个基本只是是单警号#变成字符串，双井号##黏贴左右。

那么从下面开始看起，首先声明并了一个函数，举个例子，此处的type为Concat，那么函数名字是Creator_ConcatLayer。这个函数就是把param传到了类ConcatLayer里面new出一个实例，然后返回对应的smart pointer。

然后进入到上面的宏。声明了两个静态全局变量，这两个静态全局变量都是LayerRegisterer类型的，前面已经说过，这个类型除了实例化的时候起点作用，其他时候不齐作用，所以变量本身不重要，重要的是把字符串“Concat”（也就是#type）和模板实例化后的函数指针creator<float>作为类参传给LayerRegisterer，也就能传给LayerRegistry里面的字典拉，也就能注册啦。

嚯，好家伙！绕了这么一大圈，终于给盘明白了。其实它的逻辑和我写的那个简陋的python代码并无二致，就是没字典建字典，当你新写了一个类，你把这个类名字和类本身，当成一条记录插入到字典里面去，后面别人就好查询了呗~

至于solver_factory完全是一模一样的东西。而caffe.cpp也就是定义了一大堆函数，定义完每个函数后，将函数名和函数指针给存到字典里面去，需要时再查询，再用。

四、使用流程

最后再简单地聊一聊这个注册机制的使用，仍以ConcatLayer为例（src/caffe/layers/concat_layer.cpp），定义完Layer类的各个方法以后，会执行下面2行代码：

INSTANTIATE_CLASS(ConcatLayer);
REGISTER_LAYER_CLASS(Concat);

其中，宏INSTANTIATE_CLASS定义在：/include/caffe/common.hpp，

// Instantiate a class with float and double specifications.
#define INSTANTIATE_CLASS(classname) \
  char gInstantiationGuard##classname; \
  template class classname<float>; \
  template class classname<double>

定义了一个字符（没明白这一块），然后用float和double实例化了两个模板类，在这个例子里面，我们把ConcatLayer<Dtype>给实例化成ConcatLayer<float>和ConcatLayer<double>。

那么第2行的REGISTER_LAYER_CLASS在前文已经分析的比较清楚了，我们把“Concat”和对应的返回ConcatLayer类实例的函数，添加到字典中。

最后，在Net的初始化里面（src/caffe/net.cpp），

template <typename Dtype>
void Net<Dtype>::Init(const NetParameter& in_param) {
    // ... ...
    for (int layer_id = 0; layer_id < param.layer_size(); ++layer_id) {
        // ... ...
        // Setup layer.
        const LayerParameter& layer_param = param.layer(layer_id);
        layers_.push_back(LayerRegistry<Dtype>::CreateLayer(layer_param));  // vector<shared_ptr<Layer<Dtype> > > layers_;
        layers_[layer_id]->SetUp(bottom_vecs_[layer_id], top_vecs_[layer_id]);

有选择地截取了net.cpp里面的一些东西，可以看到，在Init方法里面，通过对layer_id的for循环，我们拿到对应的layer_param，然后传到了CreateLayer里面去创建Layer类实例的smart pointer，一并存储到layer_这个vector里面，后面就拿layers_这个数据结构做setup，forward，backward等等一系列事啦。

五、总结

简单地看了一下caffe中的注册机制，其实很多时候读不懂代码，一方面固然是基础薄弱，但更重要的还是不知道他想去干什么。如果能把顶层的逻辑捋清，相信一切都能够迎刃而解~~