1. 程式人生 > >caffe原始碼解析:層(layer)的註冊與管理

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_ConvolutionLayer;

如果我們把最終所有類的註冊展開的話,會是下面的形式(注意這個是對應資料型別為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的建設工作。可以看到,這個CreateLayer先是得到型別type=“AbsVal”,然後得到registry也就該類對應的註冊資訊,最後通過registry[type](param)例項化AbsVal這個類。

另外,值得注意的一點是,像Convolution這樣不能直接自己構造的類(因為有GPU和CPU模式),註冊的建構函式有所不同,如

REGISTER_LAYER_CREATOR(Convolution, GetConvolutionLayer);