Merge pull request #2930 from lukeyeager/pycaffe-layer_type_list

Expose LayerFactory::LayerTypeList in pycaffe

3 files changed, 13 insertions, 1 deletions

diff --git a/python/caffe/__init__.py b/python/caffe/__init__.py
index 1b2da510..6cc44e72 100644
--- a/python/caffe/__init__.py
+++ b/python/caffe/__init__.py
@@ -1,5 +1,5 @@
 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
diff --git a/python/caffe/_caffe.cpp b/python/caffe/_caffe.cpp
index e1ae3ec7..020a5bee 100644
--- a/python/caffe/_caffe.cpp
+++ b/python/caffe/_caffe.cpp
@@ -200,6 +200,8 @@ BOOST_PYTHON_MODULE(_caffe) {
   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))
diff --git a/python/caffe/test/test_layer_type_list.py b/python/caffe/test/test_layer_type_list.py
new file mode 100644
index 00000000..7edc80df
--- /dev/null
+++ b/python/caffe/test/test_layer_type_list.py
@@ -0,0 +1,10 @@
+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)