include/caffe/layer_factory.hpp

   1 /**
   2  * @brief A layer factory that allows one to register layers.
   3  * During runtime, registered layers can be called by passing a LayerParameter
   4  * protobuffer to the CreateLayer function:
   5  *
   6  *     LayerRegistry<Dtype>::CreateLayer(param);
   7  *
   8  * There are two ways to register a layer. Assuming that we have a layer like:
   9  *
  10  *   template <typename Dtype>
  11  *   class MyAwesomeLayer : public Layer<Dtype> {
  12  *     // your implementations
  13  *   };
  14  *
  15  * and its type is its C++ class name, but without the "Layer" at the end
  16  * ("MyAwesomeLayer" -> "MyAwesome").
  17  *
  18  * If the layer is going to be created simply by its constructor, in your c++
  19  * file, add the following line:
  20  *
  21  *    REGISTER_LAYER_CLASS(MyAwesome);
  22  *
  23  * Or, if the layer is going to be created by another creator function, in the
  24  * format of:
  25  *
  26  *    template <typename Dtype>
  27  *    Layer<Dtype*> GetMyAwesomeLayer(const LayerParameter& param) {
  28  *      // your implementation
  29  *    }
  30  *
  31  * (for example, when your layer has multiple backends, see GetConvolutionLayer
  32  * for a use case), then you can register the creator function instead, like
  33  *
  34  * REGISTER_LAYER_CREATOR(MyAwesome, GetMyAwesomeLayer)
  35  *
  36  * Note that each layer type should only be registered once.
  37  */
  38
  39 #ifndef CAFFE_LAYER_FACTORY_H_
  40 #define CAFFE_LAYER_FACTORY_H_
  41
  42 #include <map>
  43 #include <string>
  44 #include <vector>
  45
  46 #include "caffe/common.hpp"
  47 #include "caffe/layer.hpp"
  48 #include "caffe/proto/caffe.pb.h"
  49
  50 namespace caffe {
  51
  52 template <typename Dtype>
  53 class Layer;
  54
  55 template <typename Dtype>
  56 class LayerRegistry {
  57  public:
  58   typedef shared_ptr<Layer<Dtype> > (*Creator)(const LayerParameter&);
  59   typedef std::map<string, Creator> CreatorRegistry;
  60
  61   static CreatorRegistry& Registry() {
  62     static CreatorRegistry* g_registry_ = new CreatorRegistry();
  63     return *g_registry_;
  64   }
  65
  66   // Adds a creator.
  67   static void AddCreator(const string& type, Creator creator) {
  68     CreatorRegistry& registry = Registry();
  69     CHECK_EQ(registry.count(type), 0)
  70         << "Layer type " << type << " already registered.";
  71     registry[type] = creator;
  72   }
  73
  74   // Get a layer using a LayerParameter.
  75   static shared_ptr<Layer<Dtype> > CreateLayer(const LayerParameter& param) {
  76     if (Caffe::root_solver()) {
  77       LOG(INFO) << "Creating layer " << param.name();
  78     }
  79     const string& type = param.type();
  80     CreatorRegistry& registry = Registry();
  81     CHECK_EQ(registry.count(type), 1) << "Unknown layer type: " << type
  82         << " (known types: " << LayerTypeListString() << ")";
  83     return registry[type](param);
  84   }
  85
  86   static vector<string> LayerTypeList() {
  87     CreatorRegistry& registry = Registry();
  88     vector<string> layer_types;
  89     for (typename CreatorRegistry::iterator iter = registry.begin();
  90          iter != registry.end(); ++iter) {
  91       layer_types.push_back(iter->first);
  92     }
  93     return layer_types;
  94   }
  95
  96  private:
  97   // Layer registry should never be instantiated - everything is done with its
  98   // static variables.
  99   LayerRegistry() {}
 100
 101   static string LayerTypeListString() {
 102     vector<string> layer_types = LayerTypeList();
 103     string layer_types_str;
 104     for (vector<string>::iterator iter = layer_types.begin();
 105          iter != layer_types.end(); ++iter) {
 106       if (iter != layer_types.begin()) {
 107         layer_types_str += ", ";
 108       }
 109       layer_types_str += *iter;
 110     }
 111     return layer_types_str;
 112   }
 113 };
 114
 115
 116 template <typename Dtype>
 117 class LayerRegisterer {
 118  public:
 119   LayerRegisterer(const string& type,
 120                   shared_ptr<Layer<Dtype> > (*creator)(const LayerParameter&)) {
 121     // LOG(INFO) << "Registering layer type: " << type;
 122     LayerRegistry<Dtype>::AddCreator(type, creator);
 123   }
 124 };
 125
 126
 127 #define REGISTER_LAYER_CREATOR(type, creator)                                  \
 128   static LayerRegisterer<float> g_creator_f_##type(#type, creator<float>);     \
 129   static LayerRegisterer<double> g_creator_d_##type(#type, creator<double>)    \
 130
 131 #define REGISTER_LAYER_CLASS(type)                                             \
 132   template <typename Dtype>                                                    \
 133   shared_ptr<Layer<Dtype> > Creator_##type##Layer(const LayerParameter& param) \
 134   {                                                                            \
 135     return shared_ptr<Layer<Dtype> >(new type##Layer<Dtype>(param));           \
 136   }                                                                            \
 137   REGISTER_LAYER_CREATOR(type, Creator_##type##Layer)
 138
 139 }  // namespace caffe
 140
 141 #endif  // CAFFE_LAYER_FACTORY_H_