#include <map>
#include <string>
+#include <vector>
#include "caffe/common.hpp"
#include "caffe/proto/caffe.pb.h"
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;
}
};
from .pycaffe import Net, SGDSolver
-from ._caffe import set_mode_cpu, set_mode_gpu, set_device, Layer, get_solver
+from ._caffe import set_mode_cpu, set_mode_gpu, set_device, Layer, get_solver, layer_type_list
from .proto.caffe_pb2 import TRAIN, TEST
from .classifier import Classifier
from .detector import Detector
bp::def("set_mode_gpu", &set_mode_gpu);
bp::def("set_device", &Caffe::SetDevice);
+ bp::def("layer_type_list", &LayerRegistry<Dtype>::LayerTypeList);
+
bp::class_<Net<Dtype>, shared_ptr<Net<Dtype> >, boost::noncopyable >("Net",
bp::no_init)
.def("__init__", bp::make_constructor(&Net_Init))
--- /dev/null
+import unittest
+
+import caffe
+
+class TestLayerTypeList(unittest.TestCase):
+
+ def test_standard_types(self):
+ for type_name in ['Data', 'Convolution', 'InnerProduct']:
+ self.assertIn(type_name, caffe.layer_type_list(),
+ '%s not in layer_type_list()' % type_name)