Caffe-HRT/include/caffe/layers/python_layer.hpp at master · OAID/Caffe-HRT · GitHub

Repository navigation

History

59 lines (47 loc) · 1.73 KB

Raw

#ifndef CAFFE_PYTHON_LAYER_HPP_

#define CAFFE_PYTHON_LAYER_HPP_

#include <boost/python.hpp>

#include <vector>

#include "caffe/layer.hpp"

namespace bp = boost::python;

namespace caffe {

template <typename Dtype>

class PythonLayer : public Layer<Dtype> {

public:

PythonLayer(PyObject* self, const LayerParameter& param)

: Layer<Dtype>(param), self_(bp::handle<>(bp::borrowed(self))) { }

virtual void LayerSetUp(const vector<Blob<Dtype>*>& bottom,

const vector<Blob<Dtype>*>& top) {

// Disallow PythonLayer in MultiGPU training stage, due to GIL issues

// Details: https://github.com/BVLC/caffe/issues/2936

if (this->phase_ == TRAIN && Caffe::solver_count() > 1

&& !Caffe::multiprocess()) {

LOG(FATAL) << "PythonLayer does not support CLI Multi-GPU, use train.py";

}

self_.attr("param_str") = bp::str(

this->layer_param_.python_param().param_str());

self_.attr("phase") = static_cast<int>(this->phase_);

self_.attr("setup")(bottom, top);

}

virtual void Reshape(const vector<Blob<Dtype>*>& bottom,

const vector<Blob<Dtype>*>& top) {

self_.attr("reshape")(bottom, top);

}

virtual inline bool ShareInParallel() const {

return this->layer_param_.python_param().share_in_parallel();

}

virtual inline const char* type() const { return "Python"; }

protected:

virtual void Forward_cpu(const vector<Blob<Dtype>*>& bottom,

const vector<Blob<Dtype>*>& top) {

self_.attr("forward")(bottom, top);

}

virtual void Backward_cpu(const vector<Blob<Dtype>*>& top,

const vector<bool>& propagate_down, const vector<Blob<Dtype>*>& bottom) {

self_.attr("backward")(top, propagate_down, bottom);

}

private:

bp::object self_;

};

} // namespace caffe

#endif

Caffe-HRT/include/caffe/layers/python_layer.hpp at master · OAID/Caffe-HRT · GitHub

Repository navigation

Expand file tree

Breadcrumbs

python_layer.hpp

Latest commit

History

Breadcrumbs

python_layer.hpp

File metadata and controls