Python layer for rapidly writing nets in Python

Makefile

@@ -103,6 +103,9 @@ OBJ_BUILD_DIR := $(BUILD_DIR)/src/$(PROJECT)
 LAYER_BUILD_DIR := $(OBJ_BUILD_DIR)/layers
 UTIL_BUILD_DIR := $(OBJ_BUILD_DIR)/util
 OBJS := $(PROTO_OBJS) $(CXX_OBJS) $(CU_OBJS)
+ifeq ($(USE_PYTHON_LAYER), 1)
+	OBJS += python/$(PROJECT)/_$(PROJECT).o
+endif
 # tool, example, and test objects
 TOOL_OBJS := $(addprefix $(BUILD_DIR)/, ${TOOL_SRCS:.cpp=.o})
 TOOL_BUILD_DIR := $(BUILD_DIR)/tools
@@ -172,6 +175,9 @@ LIBRARIES += pthread \
 	hdf5_hl hdf5 \
 	opencv_core opencv_highgui opencv_imgproc
 PYTHON_LIBRARIES := boost_python python2.7
+ifeq ($(USE_PYTHON_LAYER), 1)
+	LIBRARIES += $(PYTHON_LIBRARIES)
+endif
 WARNINGS := -Wall -Wno-sign-compare
 
 ##############################
@@ -309,6 +315,10 @@ endif
 INCLUDE_DIRS += $(BLAS_INCLUDE)
 LIBRARY_DIRS += $(BLAS_LIB)
 
+ifeq ($(USE_PYTHON_LAYER), 1)
+	COMMON_FLAGS += -DUSE_PYTHON_LAYER
+endif
+
 # Complete build flags.
 COMMON_FLAGS += $(foreach includedir,$(INCLUDE_DIRS),-I$(includedir))
 CXXFLAGS += -pthread -fPIC $(COMMON_FLAGS) $(WARNINGS)

Makefile.config.example

@@ -7,6 +7,9 @@
 # CPU-only switch (uncomment to build without GPU support).
 # CPU_ONLY := 1
 
+# Uncomment to include the Python layer (will link caffe against Python libs).
+# USE_PYTHON_LAYER := 1
+
 # To customize your choice of compiler, uncomment and set the following.
 # N.B. the default for Linux is g++ and the default for OSX is clang++
 # CUSTOM_CXX := g++

include/caffe/python_layer.hpp

@@ -0,0 +1,51 @@
+#ifndef CAFFE_PYTHON_LAYER_HPP_
+#define CAFFE_PYTHON_LAYER_HPP_
+
+#include <boost/python.hpp>
+#include <vector>
+
+#include "../python/caffe/_caffe.hpp"
+#include "caffe/layer.hpp"
+
+namespace caffe {
+
+/**
+ * @brief Wrap a layer implemented in Python.
+ */
+template <typename Dtype>
+class PythonLayer : public Layer<Dtype> {
+ public:
+  /**
+   * @param param provides python_param, with required parameters:
+   * - module. The module to import with the layer implementation. Note that
+   *   the current directory is not in the module search path by default.
+   * - layer. The name of the layer class, which must implement setup
+   *   (for LayerSetUp), reshape (for Reshape), forward (for Forward_cpu), and
+   *   backward (for Backward_cpu).
+   */
+  explicit PythonLayer(const LayerParameter& param)
+      : Layer<Dtype>(param) {}
+  virtual void LayerSetUp(const vector<Blob<Dtype>*>& bottom,
+      const vector<Blob<Dtype>*>& top);
+  virtual void Reshape(const vector<Blob<Dtype>*>& bottom,
+      const vector<Blob<Dtype>*>& top);
+
+  virtual inline LayerParameter_LayerType type() const {
+    return LayerParameter_LayerType_PYTHON;
+  }
+
+ protected:
+  virtual void Forward_cpu(const vector<Blob<Dtype>*>& bottom,
+      const vector<Blob<Dtype>*>& top);
+  virtual void Backward_cpu(const vector<Blob<Dtype>*>& top,
+      const vector<bool>& propagate_down, const vector<Blob<Dtype>*>& bottom);
+
+  boost::python::object layer_;
+
+ private:
+  vector<PyBlob<Dtype> > PythonBlobVector(const vector<Blob<Dtype>*>& vec);
+};
+
+}  // namespace caffe
+
+#endif

python/caffe/_caffe.cpp

@@ -37,7 +37,7 @@ static void CheckFile(const string& filename) {
 bp::object PyBlobWrap::get_data() {
   npy_intp dims[] = {num(), channels(), height(), width()};
 
-  PyObject *obj = PyArray_SimpleNewFromData(4, dims, NPY_FLOAT32,
+  PyObject* obj = PyArray_SimpleNewFromData(4, dims, NPY_FLOAT32,
                                             blob_->mutable_cpu_data());
   PyArray_SetBaseObject(reinterpret_cast<PyArrayObject *>(obj), self_);
   Py_INCREF(self_);
@@ -49,9 +49,9 @@ bp::object PyBlobWrap::get_data() {
 bp::object PyBlobWrap::get_diff() {
   npy_intp dims[] = {num(), channels(), height(), width()};
 
-  PyObject *obj = PyArray_SimpleNewFromData(4, dims, NPY_FLOAT32,
+  PyObject* obj = PyArray_SimpleNewFromData(4, dims, NPY_FLOAT32,
                                             blob_->mutable_cpu_diff());
-  PyArray_SetBaseObject(reinterpret_cast<PyArrayObject *>(obj), self_);
+  PyArray_SetBaseObject(reinterpret_cast<PyArrayObject*>(obj), self_);
   Py_INCREF(self_);
   bp::handle<> h(obj);
 
@@ -198,6 +198,9 @@ BOOST_PYTHON_MODULE(_caffe) {
   bp::class_<vector<string> >("StringVec")
       .def(bp::vector_indexing_suite<vector<string> >());
 
+  bp::class_<vector<bool> >("BoolVec")
+      .def(bp::vector_indexing_suite<vector<bool> >());
+
   import_array();
 }
 

python/caffe/_caffe.hpp

@@ -18,13 +18,21 @@ using boost::shared_ptr;
 
 namespace caffe {
 
+
 // wrap shared_ptr<Blob> in a class that we construct in C++ and pass
 // to Python
 template <typename Dtype>
 class PyBlob {
  public:
-  explicit PyBlob(const shared_ptr<Blob<Dtype> > &blob)
-      : blob_(blob) {}
+  // Construct from shared_ptr: memory will be correctly managed,
+  // even if Python holds onto a Blob beyond the life of its Net.
+  explicit PyBlob(const shared_ptr<Blob<Dtype> >& blob)
+      : blob_(blob) { }
+  // Construct from raw pointer: memory will become invalid once the
+  // owning Net is deleted. This exists only so that the raw Blob*s
+  // used in the layer interface can be passed to embedded Python.
+  explicit PyBlob(Blob<Dtype>* blob)
+      : blob_(blob, null_deleter()) { }
 
   int num() const { return blob_->num(); }
   int channels() const { return blob_->channels(); }
@@ -42,6 +50,13 @@ class PyBlob {
 
  protected:
   shared_ptr<Blob<Dtype> > blob_;
+
+ private:
+  // A dummy class that lets us use raw pointers as shared_ptrs to get
+  // around the fact that layers take around raw pointers.
+  struct null_deleter {
+    void operator()(void const*) const { }
+  };
 };
 
 // We need another wrapper (used as boost::python's HeldType) that receives a
@@ -91,7 +106,6 @@ class PyNet {
 
   void Init(string param_file);
 
-
   // Generate Python exceptions for badly shaped or discontiguous arrays.
   inline void check_contiguous_array(PyArrayObject* arr, string name,
       int channels, int height, int width);

src/caffe/layer_factory.cpp

@@ -3,6 +3,9 @@
 #include "caffe/layer.hpp"
 #include "caffe/proto/caffe.pb.h"
 #include "caffe/vision_layers.hpp"
+#ifdef USE_PYTHON_LAYER
+#include "caffe/python_layer.hpp"
+#endif
 
 namespace caffe {
 
@@ -231,6 +234,13 @@ Layer<Dtype>* GetLayer(const LayerParameter& param) {
     return GetPoolingLayer<Dtype>(name, param);
   case LayerParameter_LayerType_POWER:
     return new PowerLayer<Dtype>(param);
+  case LayerParameter_LayerType_PYTHON:
+#ifdef USE_PYTHON_LAYER
+    return new PythonLayer<Dtype>(param);
+#else
+    LOG(FATAL) << "Attempt to use PythonLayer, but built without "
+                  "USE_PYTHON_LAYER option.";
+#endif
   case LayerParameter_LayerType_RELU:
     return GetReLULayer<Dtype>(name, param);
   case LayerParameter_LayerType_SILENCE:

src/caffe/layers/python_layer.cpp

@@ -0,0 +1,74 @@
+#ifdef USE_PYTHON_LAYER
+#include <boost/python.hpp>
+#include <Python.h>
+#include <vector>
+
+#include "caffe/layer.hpp"
+#include "caffe/python_layer.hpp"
+
+namespace bp = boost::python;
+
+namespace caffe {
+
+template <typename Dtype>
+vector<PyBlob<Dtype> > PythonLayer<Dtype>::PythonBlobVector(
+    const vector<Blob<Dtype>*>& vec) {
+  return vector<PyBlob<Dtype> >(vec.begin(), vec.end());
+}
+
+template <typename Dtype>
+void PythonLayer<Dtype>::LayerSetUp(const vector<Blob<Dtype>*>& bottom,
+     const vector<Blob<Dtype>*>& top) {
+  Py_Initialize();
+  init_caffe();
+
+  try {
+    bp::object module_ = bp::import(
+        this->layer_param_.python_param().module().c_str());
+    layer_ = module_.attr(this->layer_param_.python_param().layer().c_str())();
+
+    layer_.attr("setup")(PythonBlobVector(bottom), PythonBlobVector(top));
+  } catch (bp::error_already_set) {
+    PyErr_Print();
+    throw;
+  }
+}
+
+template <typename Dtype>
+void PythonLayer<Dtype>::Reshape(const vector<Blob<Dtype>*>& bottom,
+      const vector<Blob<Dtype>*>& top) {
+  try {
+    layer_.attr("reshape")(PythonBlobVector(bottom), PythonBlobVector(top));
+  } catch (bp::error_already_set) {
+    PyErr_Print();
+    throw;
+  }
+}
+
+template <typename Dtype>
+void PythonLayer<Dtype>::Forward_cpu(const vector<Blob<Dtype>*>& bottom,
+      const vector<Blob<Dtype>*>& top) {
+  try {
+    layer_.attr("forward")(PythonBlobVector(bottom), PythonBlobVector(top));
+  } catch (bp::error_already_set) {
+    PyErr_Print();
+    throw;
+  }
+}
+
+template <typename Dtype>
+void PythonLayer<Dtype>::Backward_cpu(const vector<Blob<Dtype>*>& top,
+    const vector<bool>& propagate_down, const vector<Blob<Dtype>*>& bottom) {
+  try {
+    layer_.attr("backward")(PythonBlobVector(top), propagate_down,
+        PythonBlobVector(bottom));
+  } catch (bp::error_already_set) {
+    PyErr_Print();
+    throw;
+  }
+}
+
+INSTANTIATE_CLASS(PythonLayer);
+
+}  // namespace caffe
+#endif

src/caffe/proto/caffe.proto

@@ -198,7 +198,7 @@ message NetStateRule {
 // NOTE
 // Update the next available ID when you add a new LayerParameter field.
 //
-// LayerParameter next available ID: 41 (last added: contrastive_loss_param)
+// LayerParameter next available ID: 42 (last added: python_param)
 message LayerParameter {
   repeated string bottom = 2; // the name of the bottom blobs
   repeated string top = 3; // the name of the top blobs
@@ -219,7 +219,7 @@ message LayerParameter {
   // line above the enum. Update the next available ID when you add a new
   // LayerType.
   //
-  // LayerType next available ID: 38 (last added: CONTRASTIVE_LOSS)
+  // LayerType next available ID: 39 (last added: PYTHON)
   enum LayerType {
     // "NONE" layer type is 0th enum element so that we don't cause confusion
     // by defaulting to an existent LayerType (instead, should usually error if
@@ -251,6 +251,7 @@ message LayerParameter {
     MVN = 34;
     POOLING = 17;
     POWER = 26;
+    PYTHON = 38;
     RELU = 18;
     SIGMOID = 19;
     SIGMOID_CROSS_ENTROPY_LOSS = 27;
@@ -310,6 +311,7 @@ message LayerParameter {
   optional MVNParameter mvn_param = 34;
   optional PoolingParameter pooling_param = 19;
   optional PowerParameter power_param = 21;
+  optional PythonParameter python_param = 41;
   optional ReLUParameter relu_param = 30;
   optional SigmoidParameter sigmoid_param = 38;
   optional SoftmaxParameter softmax_param = 39;
@@ -603,6 +605,12 @@ message PowerParameter {
   optional float shift = 3 [default = 0.0];
 }
 
+// Message that stores parameters used by PythonLayer
+message PythonParameter {
+  optional string module = 1;
+  optional string layer = 2;
+}
+
 // Message that stores parameters used by ReLULayer
 message ReLUParameter {
   // Allow non-zero slope for negative inputs to speed up optimization