1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
|
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])
|
