|
|
|
|
|
|
|
|
|
|
| import unittest |
| import numpy as np |
| from caffe2.python import workspace, core |
|
|
| from caffe2.proto import caffe2_pb2 |
|
|
|
|
| class TestPredictor(unittest.TestCase): |
| def setUp(self): |
| np.random.seed(1) |
| self.predict_net = self._predict_net |
| self.init_net = self._init_net |
|
|
| @property |
| def _predict_net(self): |
| net = caffe2_pb2.NetDef() |
| net.name = 'test-predict-net' |
| net.external_input[:] = ['A', 'B'] |
| net.external_output[:] = ['C'] |
| net.op.extend([ |
| core.CreateOperator( |
| 'MatMul', |
| ['A', 'B'], |
| ['C'], |
| ) |
| ]) |
| return net.SerializeToString() |
|
|
| @property |
| def _init_net(self): |
| net = caffe2_pb2.NetDef() |
| net.name = 'test-init-net' |
| net.external_output[:] = ['A', 'B'] |
| net.op.extend([ |
| core.CreateOperator( |
| 'GivenTensorFill', |
| [], |
| ['A'], |
| shape=(2, 3), |
| values=np.zeros((2, 3), np.float32).flatten().tolist(), |
| ), |
| core.CreateOperator( |
| 'GivenTensorFill', |
| [], |
| ['B'], |
| shape=(3, 4), |
| values=np.zeros((3, 4), np.float32).flatten().tolist(), |
| ), |
| ]) |
| return net.SerializeToString() |
|
|
| def test_run(self): |
| A = np.ones((2, 3), np.float32) |
| B = np.ones((3, 4), np.float32) |
| predictor = workspace.Predictor(self.init_net, self.predict_net) |
| outputs = predictor.run([A, B]) |
| self.assertEqual(len(outputs), 1) |
| np.testing.assert_almost_equal(np.dot(A, B), outputs[0]) |
|
|
| def test_run_map(self): |
| A = np.zeros((2, 3), np.float32) |
| B = np.ones((3, 4), np.float32) |
| predictor = workspace.Predictor(self.init_net, self.predict_net) |
| outputs = predictor.run({ |
| 'B': B, |
| }) |
| self.assertEqual(len(outputs), 1) |
| np.testing.assert_almost_equal(np.dot(A, B), outputs[0]) |
|
|
