diff options
Diffstat (limited to 'include/caffe')
-rw-r--r-- | include/caffe/layer_factory.hpp | 31 |
1 files changed, 21 insertions, 10 deletions
diff --git a/include/caffe/layer_factory.hpp b/include/caffe/layer_factory.hpp index 32e849de..2c2fde4d 100644 --- a/include/caffe/layer_factory.hpp +++ b/include/caffe/layer_factory.hpp @@ -41,6 +41,7 @@ #include <map> #include <string> +#include <vector> #include "caffe/common.hpp" #include "caffe/proto/caffe.pb.h" @@ -77,26 +78,36 @@ class LayerRegistry { const string& type = param.type(); CreatorRegistry& registry = Registry(); CHECK_EQ(registry.count(type), 1) << "Unknown layer type: " << type - << " (known types: " << LayerTypeList() << ")"; + << " (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 LayerTypeList() { - CreatorRegistry& registry = Registry(); - string layer_types; - for (typename CreatorRegistry::iterator iter = registry.begin(); - iter != registry.end(); ++iter) { - if (iter != registry.begin()) { - layer_types += ", "; + 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 += iter->first; + layer_types_str += *iter; } - return layer_types; + return layer_types_str; } }; |