From 56f324d1912d55d780232c6e47d75ee48f39ae2c Mon Sep 17 00:00:00 2001 From: Alexander Sidorov Date: Wed, 15 Mar 2017 11:08:53 -0700 Subject: [PATCH] Added predictor bindings to python interface Summary: from caffe2.python import workspace; p = workspace.Predictor(init_net, predict_net); outputs = p.run(inputs) Reviewed By: Yangqing Differential Revision: D4576793 fbshipit-source-id: b829bbcaf2e7c34dad85024177433207bd96a234 --- caffe2/python/pybind_state.cc | 29 +++++++++++++- caffe2/python/workspace.py | 1 + caffe2/python/workspace_test.py | 46 +++++++++++++++++++++- caffe2/test/assets/squeeze_predict_net.pb | Bin 0 -> 6175 bytes 4 files changed, 73 insertions(+), 3 deletions(-) create mode 100644 caffe2/test/assets/squeeze_predict_net.pb diff --git a/caffe2/python/pybind_state.cc b/caffe2/python/pybind_state.cc index ce7ec4019af..633c71efc59 100644 --- a/caffe2/python/pybind_state.cc +++ b/caffe2/python/pybind_state.cc @@ -536,11 +536,38 @@ void addObjectMethods(py::module& m) { .def( "__init__", [](Predictor& instance, py::bytes init_net, py::bytes predict_net) { + CAFFE_ENFORCE(gWorkspace); NetDef init_net_, predict_net_; CAFFE_ENFORCE(ParseProtobufFromLargeString(init_net, &init_net_)); CAFFE_ENFORCE( ParseProtobufFromLargeString(predict_net, &predict_net_)); - new (&instance) Predictor(init_net_, predict_net_); + new (&instance) Predictor(init_net_, predict_net_, gWorkspace); + }) + .def( + "run", + [](Predictor& instance, + std::vector inputs) -> std::vector { + std::vector tensors; + std::vector tensors_data(inputs.size()); + for (auto i = 0; i < inputs.size(); ++i) { + auto input = inputs[i]; + CAFFE_ENFORCE( + PyArray_Check(input.ptr()), + "Input must be of type numpy array."); + PyArrayObject* array = + reinterpret_cast(input.ptr()); + TensorFeeder().FeedTensor( + DeviceOption(), array, &(tensors_data[i])); + tensors.push_back(&(tensors_data[i])); + } + std::vector out; + instance.run(tensors, &out); + std::vector pyout; + for (auto t : out) { + pyout.push_back( + TensorFetcher().FetchTensor(*t, true).obj); + } + return pyout; }); } diff --git a/caffe2/python/workspace.py b/caffe2/python/workspace.py index 0fdcce9335f..b8515621a82 100644 --- a/caffe2/python/workspace.py +++ b/caffe2/python/workspace.py @@ -27,6 +27,7 @@ SwitchWorkspace = C.switch_workspace RootFolder = C.root_folder Workspaces = C.workspaces BenchmarkNet = C.benchmark_net +Predictor = C.Predictor is_asan = C.is_asan has_gpu_support = C.has_gpu_support diff --git a/caffe2/python/workspace_test.py b/caffe2/python/workspace_test.py index 3ccb496f6e6..eebc13eb739 100644 --- a/caffe2/python/workspace_test.py +++ b/caffe2/python/workspace_test.py @@ -3,12 +3,11 @@ import os import unittest from caffe2.proto import caffe2_pb2 -from caffe2.python import core, test_util, workspace +from caffe2.python import core, test_util, workspace, cnn import caffe2.python.hypothesis_test_util as htu import hypothesis.strategies as st from hypothesis import given -from caffe2.proto import caffe2_pb2 class TestWorkspace(unittest.TestCase): @@ -469,5 +468,48 @@ class TestCWorkspace(htu.HypothesisTestCase): ws.create_net("...") +class TestPredictor(unittest.TestCase): + def _create_model(self): + m = cnn.CNNModelHelper() + y = m.FC("data", "y", + dim_in=4, dim_out=2, + weight_init=m.ConstantInit(1.0), + bias_init=m.ConstantInit(0.0), + axis=0) + m.net.AddExternalOutput(y) + return m + + # Use this test with a bigger model to see how using Predictor allows to + # avoid issues with low protobuf size limit in Python + # + # def test_predictor_predefined(self): + # workspace.ResetWorkspace() + # path = 'caffe2/caffe2/test/assets/' + # with open(path + 'squeeze_predict_net.pb') as f: + # self.predict_net = f.read() + # with open(path + 'squeeze_init_net.pb') as f: + # self.init_net = f.read() + # self.predictor = workspace.Predictor(self.init_net, self.predict_net) + + # inputs = [np.zeros((1, 3, 256, 256), dtype='f')] + # outputs = self.predictor.run(inputs) + # self.assertEqual(len(outputs), 1) + # self.assertEqual(outputs[0].shape, (1, 1000, 1, 1)) + # self.assertAlmostEqual(outputs[0][0][0][0][0], 5.19026289e-05) + + + def test_predictor_memory_model(self): + workspace.ResetWorkspace() + m = self._create_model() + workspace.FeedBlob("data", np.zeros([4], dtype='float32')) + self.predictor = workspace.Predictor( + workspace.StringifyProto(m.param_init_net.Proto()), + workspace.StringifyProto(m.net.Proto())) + + inputs = np.array([1, 3, 256, 256], dtype='float32') + outputs = self.predictor.run([inputs]) + np.testing.assert_array_almost_equal(np.array([[516, 516]], dtype='float32'), outputs) + + if __name__ == '__main__': unittest.main() diff --git a/caffe2/test/assets/squeeze_predict_net.pb b/caffe2/test/assets/squeeze_predict_net.pb new file mode 100644 index 0000000000000000000000000000000000000000..a06d95947321a407aa9b4c1d0c00d320af80f426 GIT binary patch literal 6175 zcmb`L%Wu;_5XMvb2wL7NK)rCu3FT5snm%$tJcPJ_cu5?Ph@8fhsA(ONlm`BKc6VmX zH)|6%5;yzvHxGM$+1P6u_NCbvX7kK6kBjpkyZuW!7}(9Av^Ji*3vFYzX7lN2m~UTc zH_iGa8*Vqec$!bo^YQk&y=%x{hP}_n7xq)BPKsi@+iJef7H{3rKl8P3n9X83%%`o5 zH_u;vYTqZfFGamM>F>R* zDI}uaT%_(Tqzb)#uH1zdeCTz0LJrE~Y;J8IPy55s*{rqYNOjjDbZlW%DACcWAx3FT z3~GPBZ~Q5g#9$~9gGy@yr-D~#1f87Q zB2xa=YKehly0?%Di4mv!T@f)NI42R)*-^yE=~8085HaGlVX5i-_r^H6-AVI+#scP$0$N1jT?m2-aL1VqC|Ts z>MiiI5{i&a_ZCtiisCF$6u~(O#eI#UoGwL?7ccP=F$XooD2<6hEp;YL42BXhsOYKS z6(Z)KN(@Ro6_8T9V`3nZrUF80ZypaOMxwnG^%lmw5;2fW_ZCtiG2$$V5y3f$m;+6W zoTd;XULxkOh8U$WF{q`^go(jWA_f&b6}&>k99D@ziKhZmYIjTwMAB41NbSwz!Nf?k zm!jUnm{%eOlIh+;DkMgnB{3p6ClPa~iILM3V#G_t9Muq`G$sbM)R{0b7)r#TqNjpa zh?t`)F(~m=KuYb7iGfI(3J9sac|4dHiS|;|TNv|7#6U9LTS$e(h_fU{1m`4Tjx;fH zx|EoY^mQlR<%!O+*=ei!Vp>d!i+TIH*_>wcQE{))Xgq0OGtJSgKhJ0L?KS&OXEj%k z{eIdb_XOo=V7Ew1o#lPQ0e=F+)A{~K4|`ybTGyWb$fwy!9&N3aU)i`n(TMg(;;pr| zH8)Sj#ULB^efH7$iMzW8W{XN$Gb_H%&$5Mk9IfUDI_`xn#GY&+%E6#l-H5xa(FQ%W z>Ex!}zpP4;%cB^%yu>weKgHQXA4=nFvRwc*QHoqpQ(RCJKb)ne(m0!Z7eY;xA{W#Y z7u3X0FR7_C&L&?LP!pxd1vSM5HSrx=YATJh$txt(L@9DXO>sd@d`pv>O5<$u5&|_* zid;}rTu>AL=cT67IGg;9K~0n*7t|CN)Wm-Zsi`#1CJQ0dL@9DXO>t38S(+|;ax{1b GKk(oG+h8pK literal 0 HcmV?d00001