caffe源码解析:层(layer)的注册与管理
caffe中所有的layer都是类的结构,它们的构造相关的函数都注册在一个全局变量g_registry_ 中。
首先这个变量的类型 CreatorRegistry是一个map定义,
public:
typedef shared_ptr<Layer<Dtype> > (*Creator)(const LayerParameter&);
typedef std::map<string, Creator> CreatorRegistry;
注意Creator会根据数据类型(如float,double)的不同而不同,
该定义对应的变量是g_registry_(实际是对应float有一个,对应double有一个)
层的注册函数是通过下边这个Registry() 函数实现的,
static CreatorRegistry& Registry() {
static CreatorRegistry* g_registry_ = new CreatorRegistry();
return *g_registry_;
}
注意static 这个关键字,这是一个静态结构,在caffe程序启动时(通过REGISTER_LAYER_CLASS(……)这个宏)创建g_registry_ ,(因为是static,所以只在创建时运行一次),以后各个层通过REGISTER_LAYER_CLASS(……)这个宏调用时,返回的是注册到g_registry_的具体的层。
当然我们可以看一眼这个宏定义,
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);
}
};
#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)
} // namespace caffe
比如对AbsVal层,REGISTER_LAYER_CLASS 这个宏首先定义了一个Creator_AbsValLayer函数,然后通过宏
REGISTER_LAYER_CREATOR(type, Creator_AbsValLayer)定义了一个LayerRegisterer(实际是两个LayerRegisterer,根据数据类型是float还是double,各定义了一个)静态实例,这个定义相当根据数据类型于立即执行了
AddCreateor("Convolution", Creator_AbsValLayer)这个函数,
// 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;
}
也就是把这个Creator_AbsValLayer函数注册到数组
registry["Convolution"] = Creator_AbsValLayer;
如果我们把最终所有类的注册展开的话,会是下面的形式(注意这个是对应数据类型为float的registry,对应数据类型为double的也有一个同样结构的registery)
在构造网络时,net::init函数中的for循环会根据网络的类型逐层构建,通过
layers_.push_back(LayerRegistry<Dtype>::CreateLayer(layer_param));
调用
// 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);
}
完成这个layer的建设工作。当然,取决于不同的层,调用函数是不同的,例如对convolution层,调用函数是
shared_ptr<Layer<Dtype> > GetConvolutionLayer(const LayerParameter& param)
以AbsVal为例,这个CreateLayer先是得到类型type=“AbsVal”,然后得到registry也就该类对应的注册信息,最后通过registry[type](param)实例化AbsVal这个类。
另外,值得注意的一点是,像Convolution这样不能直接自己构造的类(因为有GPU和CPU模式),注册的构造函数有所不同,如
REGISTER_LAYER_CREATOR(Convolution, GetConvolutionLayer);