diff options
| -rw-r--r-- | Makefile | 9 | ||||
| -rw-r--r-- | src/caffe/net.cpp | 16 | ||||
| -rw-r--r-- | src/pycaffe/pycaffe.cpp | 45 |
3 files changed, 60 insertions, 10 deletions
@@ -57,9 +57,14 @@ LDFLAGS += $(foreach librarydir,$(LIBRARY_DIRS),-L$(librarydir)) \ NVCC = nvcc $(NVCCFLAGS) $(CPPFLAGS) $(CUDA_ARCH) -.PHONY: all test clean distclean linecount examples distribute +.PHONY: all test clean distclean linecount examples pycaffe distribute -all: $(NAME) $(STATIC_NAME) test examples +all: $(NAME) $(STATIC_NAME) test examples pycaffe + +pycaffe: $(STATIC_NAME) src/pycaffe/pycaffe.cpp + $(CXX) -o pycaffe.so -I/usr/include/python2.7 -shared \ + src/pycaffe/pycaffe.cpp $(STATIC_NAME) $(CXXFLAGS) $(LDFLAGS) \ + $(WARNING) -lboost_python -lpython2.7 linecount: clean cloc --read-lang-def=caffe.cloc src/caffe/ diff --git a/src/caffe/net.cpp b/src/caffe/net.cpp index 3266064d..5bc67a85 100644 --- a/src/caffe/net.cpp +++ b/src/caffe/net.cpp @@ -5,8 +5,6 @@ #include <string> #include <vector> -#include <boost/python.hpp> - #include "caffe/proto/caffe.pb.h" #include "caffe/layer.hpp" #include "caffe/net.hpp" @@ -240,17 +238,19 @@ const vector<Blob<Dtype>*>& Net<Dtype>::Forward( template <typename Dtype> string Net<Dtype>::Forward(const string& input_blob_protos) { BlobProtoVector blob_proto_vec; - blob_proto_vec.ParseFromString(input_blob_protos); - CHECK_EQ(blob_proto_vec.blobs_size(), net_input_blob_indices_.size()) - << "Incorrect input size."; - for (int i = 0; i < blob_proto_vec.blobs_size(); ++i) { - blobs_[net_input_blob_indices_[i]]->FromProto(blob_proto_vec.blobs(i)); + if (net_input_blob_indices_.size()) { + blob_proto_vec.ParseFromString(input_blob_protos); + CHECK_EQ(blob_proto_vec.blobs_size(), net_input_blob_indices_.size()) + << "Incorrect input size."; + for (int i = 0; i < blob_proto_vec.blobs_size(); ++i) { + blobs_[net_input_blob_indices_[i]]->FromProto(blob_proto_vec.blobs(i)); + } } for (int i = 0; i < layers_.size(); ++i) { layers_[i]->Forward(bottom_vecs_[i], &top_vecs_[i]); } blob_proto_vec.Clear(); - for (int i = 0; i < layers_.size(); ++i) { + for (int i = 0; i < net_output_blobs_.size(); ++i) { net_output_blobs_[i]->ToProto(blob_proto_vec.add_blobs()); } string output; diff --git a/src/pycaffe/pycaffe.cpp b/src/pycaffe/pycaffe.cpp new file mode 100644 index 00000000..fb475951 --- /dev/null +++ b/src/pycaffe/pycaffe.cpp @@ -0,0 +1,45 @@ +#include <boost/python.hpp> +#include "caffe/caffe.hpp" + +using namespace caffe; +using namespace boost::python; + +// For python, we will simply use float. This is wrapped in a +#define PTYPE float + +// A simple wrapper over CaffeNet that runs the forward process. +struct CaffeNet +{ + CaffeNet(string param_file, string pretrained_param_file, + boost::python::list bottom) { + vector<int> bottom_vec; + for (int i = 0; i < boost::python::len(bottom); ++i) { + bottom_vec.push_back(boost::python::extract<int>(bottom[i])); + } + net_.reset(new Net<float>(param_file, bottom_vec)); + net_->CopyTrainedLayersFrom(pretrained_param_file); + } + + string Forward(const string& input_blobs) { + return net_->Forward(input_blobs); + } + + void set_mode_cpu() { Caffe::set_mode(Caffe::CPU); } + void set_mode_gpu() { Caffe::set_mode(Caffe::GPU); } + + void set_phase_train() { Caffe::set_phase(Caffe::TRAIN); } + void set_phase_test() { Caffe::set_phase(Caffe::TEST); } + + shared_ptr<Net<float> > net_; +}; + +BOOST_PYTHON_MODULE(pycaffe) +{ + class_<CaffeNet>("CaffeNet", init<string, string, boost::python::list>()) + .def("Forward", &CaffeNet::Forward) + .def("set_mode_cpu", &CaffeNet::set_mode_cpu) + .def("set_mode_gpu", &CaffeNet::set_mode_gpu) + .def("set_phase_train", &CaffeNet::set_phase_train) + .def("set_phase_test", &CaffeNet::set_phase_test) + ; +}
\ No newline at end of file |