pytorch

1

2

3

4

5

6
from caffe2.python import core, test_util
7
from caffe2.proto import caffe2_pb2
8
import caffe2.python.nomnigraph as ng
9

10
from hypothesis import given
11
import hypothesis.strategies as st
12
import random
13

14

15
class TestBindings(test_util.TestCase):
16
    def test_simple(self):
17
        nn = ng.NNModule()
18
        dfg = nn.dataFlow
19
        dfg.createNode(ng.NeuralNetData("X"))
20
        dfg.createNode(ng.NeuralNetOperator("FC"))
21
        assert len(nn.dataFlow.getMutableNodes()) == 2
22

23
    def test_core_net_simple(self):
24
        net = core.Net("name")
25
        net.FC(["X", "W"], ["Y"])
26
        nn = ng.NNModule(net)
27
        for node in nn.dataFlow.getMutableNodes():
28
            if node.isOperator():
29
                assert node.getName() == "FC"
30
            elif node.isTensor():
31
                assert node.getName() in ["X", "W", "Y"]
32

33
    def test_core_net_controlflow(self):
34
        net = core.Net("name")
35
        net.FC(["X", "W"], ["Y"])
36
        net.Relu(["Y"], ["Z"])
37
        nn = ng.NNModule(net)
38
        assert len(nn.controlFlow) == 2
39
        for instr in nn.controlFlow:
40
            assert instr.getType() == "Operator"
41
        assert nn.controlFlow[0].getName() == "FC"
42
        assert nn.controlFlow[1].getName() == "Relu"
43

44
    def test_core_net_nn_accessors(self):
45
        net = core.Net("name")
46
        net.FC(["X", "W"], ["Y"])
47
        net.Relu(["Y"], ["Z"])
48
        nn = ng.NNModule(net)
49
        tensors = set()
50
        for t in nn.tensors:
51
            tensors.add(t.name)
52
        assert tensors == set(["X", "W", "Y", "Z"])
53
        ops = set()
54
        for op in nn.operators:
55
            ops.add(op.name)
56
        assert ops == set(["FC", "Relu"])
57
        nodes = set()
58
        for node in nn.nodes:
59
            nodes.add(node.name)
60
        assert nodes == (ops | tensors)
61

62
    def test_netdef_simple(self):
63
        net = core.Net("name")
64
        net.FC(["X", "W"], ["Y"])
65
        nn = ng.NNModule(net.Proto())
66
        for node in nn.dataFlow.getMutableNodes():
67
            if node.isOperator():
68
                assert node.getOperator().getName() == "FC"
69
            elif node.isTensor():
70
                assert node.getTensor().getName() in ["X", "W", "Y"]
71

72
    def test_operatordef_simple(self):
73
        nn = ng.NNModule()
74
        dfg = nn.dataFlow
75
        op = core.CreateOperator("Ceil", ["X"], ["Y"], engine="CUDNN")
76
        dfg.createNode(op)
77
        for node in dfg.getMutableNodes():
78
            assert node.isOperator()
79
            assert node.getOperator().getName() == "Ceil"
80

81
    def test_invalid_node(self):
82
        nn = ng.NNModule()
83
        dfg = nn.dataFlow
84
        with self.assertRaises(Exception):
85
            dfg.createNode(7)
86

87
    def test_edges_simple(self):
88
        nn = ng.NNModule()
89
        dfg = nn.dataFlow
90
        x = dfg.createNode(ng.NeuralNetData("X"))
91
        w = dfg.createNode(ng.NeuralNetData("W"))
92
        op = dfg.createNode(ng.NeuralNetOperator("Op"))
93

94
        with self.assertRaises(Exception):
95
            dfg.createEdge(x, w)
96
        dfg.createEdge(op, w)
97
        dfg.createEdge(x, op)
98

99
        # Dot generation
100
        assert(str(dfg).startswith("digraph G"))
101

102
        # subgraph
103
        sg = ng.NNSubgraph()
104
        sg.addNode(x)
105
        sg.addNode(op)
106
        sg.induceEdges()
107
        assert len(sg) == 2
108

109
        # subgraph dot generation
110
        assert(str(sg).startswith("digraph G"))
111

112
    @given(size=st.sampled_from([10, 50]))
113
    def test_edges_complex(self, size):
114
        random.seed(1337)
115
        nn = ng.NNModule()
116
        dfg = nn.dataFlow
117

118
        data = []
119
        ops = []
120
        for _ in range(size):
121
            data.append(dfg.createNode(ng.NeuralNetData("X")))
122
        for i in range(size):
123
            ops.append(dfg.createNode(ng.NeuralNetOperator("Op" + str(i))))
124

125
        for i in range(size):
126
            for j in range(size):
127
                if bool(random.getrandbits(1)):
128
                    dfg.createEdge(data[i], ops[j])
129

130
    def test_traversal(self):
131
        net = core.Net("test")
132
        net.FC(["X", "W"], ["Y"])
133
        net.Relu(["Y"], ["Z"])
134
        nn = ng.NNModule(net)
135
        fc = nn.controlFlow[0]
136
        relu = nn.controlFlow[1]
137
        assert not fc.inputs[0].hasProducer()
138
        assert fc.inputs[0].name == "X"
139
        assert fc.inputs[1].name == "W"
140
        assert relu.outputs[0].name == "Z"
141
        assert relu.inputs[0].name == "Y"
142
        assert relu.inputs[0].hasProducer()
143
        assert relu.inputs[0].producer.name == "FC"
144
        assert fc.outputs[0].consumers[0].name == "Relu"
145

146
    def test_debug(self):
147
        nn = ng.NNModule()
148
        dfg = nn.dataFlow
149
        dfg.createNode(ng.NeuralNetData("X"))
150
        dfg.createNode(ng.NeuralNetData("W"))
151
        dfg.createNode(ng.NeuralNetOperator("Op"))
152

153
        ng.render(nn.dataFlow)
154

155
    def test_match_graph_node(self):
156
        mg = ng.NNMatchGraph()
157
        mg.createNode(ng.NeuralNetOperator("test"))
158
        nn = ng.NNModule()
159
        test = nn.dataFlow.createNode(ng.NeuralNetOperator("test"))
160
        x = nn.dataFlow.createNode(ng.NeuralNetData("X"))
161
        nn.dataFlow.createEdge(x, test)
162

163
        count = 0
164
        for match in nn.match(mg):
165
            assert len(match) == 1
166
            count += 1
167
            # Dot generation of subgraph
168
            assert(str(match).startswith("digraph G"))
169
        assert count == 1
170

171
    def test_match_graph_node_strict(self):
172
        mg = ng.NNMatchGraph()
173
        mg.createNode(ng.NeuralNetOperator("test"), strict=True)
174
        nn = ng.NNModule()
175
        test = nn.dataFlow.createNode(ng.NeuralNetOperator("test"))
176
        x = nn.dataFlow.createNode(ng.NeuralNetData("X"))
177
        nn.dataFlow.createEdge(test, x)
178

179
        count = 0
180
        for match in nn.match(mg):
181
            assert len(match) == 1
182
            count += 1
183

184
        with self.assertRaises(Exception):
185
            assert count == 1
186

187
    def test_match_graph(self):
188
        mg = ng.NNMatchGraph()
189
        test2m = mg.createNode(ng.NeuralNetOperator("test2"), strict=True)
190
        xm = mg.createNode(ng.NeuralNetData("X"), strict=True)
191
        testm = mg.createNode(ng.NeuralNetOperator("test"))
192
        mg.createEdge(test2m, xm)
193
        mg.createEdge(xm, testm)
194

195
        nn = ng.NNModule()
196
        test2 = nn.dataFlow.createNode(ng.NeuralNetOperator("test2"))
197
        x = nn.dataFlow.createNode(ng.NeuralNetData("X"))
198
        test = nn.dataFlow.createNode(ng.NeuralNetOperator("test"))
199
        nn.dataFlow.createEdge(test2, x)
200
        nn.dataFlow.createEdge(x, test)
201

202
        count = 0
203
        for match in nn.match(mg):
204
            print(len(match))
205
            assert len(match) == 3
206
            count += 1
207
        assert count == 1
208

209
    def test_delete_subgraph(self):
210
        mg = ng.NNMatchGraph()
211
        test2m = mg.createNode(ng.NeuralNetOperator("test2"), strict=True)
212
        xm = mg.createNode(ng.NeuralNetData("X"), strict=True)
213
        testm = mg.createNode(ng.NeuralNetOperator("test"))
214
        mg.createEdge(test2m, xm)
215
        mg.createEdge(xm, testm)
216

217
        nn = ng.NNModule()
218
        test2 = nn.dataFlow.createNode(ng.NeuralNetOperator("test2"))
219
        x = nn.dataFlow.createNode(ng.NeuralNetData("X"))
220
        test = nn.dataFlow.createNode(ng.NeuralNetOperator("test"))
221
        nn.dataFlow.createEdge(test2, x)
222
        nn.dataFlow.createEdge(x, test)
223

224
        for m in nn.match(mg):
225
            match = m
226
        nn.deleteSubgraph(match)
227
        assert len(nn.controlFlow) == 0
228

229
    def test_replace_subraph(self):
230
        mg = ng.NNMatchGraph()
231
        test2m = mg.createNode(ng.NeuralNetOperator("test2"), strict=True)
232
        xm = mg.createNode(ng.NeuralNetData("X"), strict=True)
233
        testm = mg.createNode(ng.NeuralNetOperator("test"))
234
        mg.createEdge(test2m, xm)
235
        mg.createEdge(xm, testm)
236

237
        nn = ng.NNModule()
238
        test2 = nn.dataFlow.createNode(ng.NeuralNetOperator("test2"))
239
        x = nn.dataFlow.createNode(ng.NeuralNetData("X"))
240
        test = nn.dataFlow.createNode(ng.NeuralNetOperator("test"))
241
        nn.dataFlow.createEdge(test2, x)
242
        nn.dataFlow.createEdge(x, test)
243

244
        for m in nn.match(mg):
245
            match = m
246
        new_op = nn.dataFlow.createNode(ng.NeuralNetOperator("new_op"))
247
        nn.replaceSubgraph(match, new_op, [], [])
248
        assert len(nn.controlFlow) == 1
249
        assert nn.controlFlow[0].name == "new_op"
250

251
    def test_genericGraph(self):
252
        g = ng.Graph()
253
        n1 = g.createNode("hello1")
254
        n2 = g.createNode("hello2")
255
        e = g.createEdge(n1, n2)
256
        ng.render(g)
257

258
    def test_createUniqueDataNode(self):
259
        net = core.Net("name")
260
        nn = ng.NNModule(net)
261
        n1 = nn.createUniqueDataNode("a")
262
        self.assertEqual(n1.name[0], "a")
263
        n2 = nn.dataFlow.createNode(ng.Operator("test1"))
264
        nn.createEdge(n1, n2)
265
        n3 = nn.createUniqueDataNode("a")
266
        nn.createEdge(n2, n3)
267
        self.assertEqual(n3.name[0], "a")
268
        self.assertNotEqual(n1.name, n3.name)
269
        n1 = nn.createUniqueDataNode("b")
270
        n2 = nn.createUniqueDataNode("b")
271
        self.assertNotEqual(n1.name, n2.name)
272

273
    def test_convertToProto(self):
274
        net = core.Net("name")
275
        net.FC(["X", "W"], ["Y"])
276
        nn = ng.NNModule(net)
277
        new_netdef = nn.convertToCaffe2Proto()
278
        print(new_netdef)
279
        print(net.Proto())
280
        assert len(new_netdef.op) == len(net.Proto().op)
281
        for i in range(len(new_netdef.op)):
282
            op = net.Proto().op[i]
283
            new_op = new_netdef.op[i]
284
            assert op.type == new_op.type
285
            assert len(op.input) == len(new_op.input)
286
            assert len(op.output) == len(new_op.output)
287
            for a, b in zip(op.input, new_op.input):
288
                assert a == b
289
            for a, b in zip(op.output, new_op.output):
290
                assert a == b
291
        for a, b in zip(new_netdef.external_input, net.Proto().external_input):
292
            assert a == b
293
        for a, b in zip(new_netdef.external_output, net.Proto().external_output):
294
            assert a == b
295

296
    def test_node_interactions(self):
297
        nn = ng.NNModule()
298
        dfg = nn.dataFlow
299
        test1 = dfg.createNode(ng.Operator("test1"))
300
        test2 = dfg.createNode(ng.Operator("test2"))
301
        x = dfg.createNode(ng.Data("x"))
302
        dfg.createEdge(test1, x)
303
        dfg.createEdge(x, test2)
304
        p = test2.getOperatorPredecessors()
305
        assert len(p) == 1
306
        assert p[0] == test1
307

308
        # Add another node
309
        test3 = dfg.createNode(ng.Operator("test3"))
310
        y = dfg.createNode(ng.Data("y"))
311
        dfg.createEdge(test3, y)
312
        dfg.createEdge(y, test2)
313
        p = test2.getOperatorPredecessors()
314
        assert len(p) == 2
315
        assert test1 in p
316
        assert test3 in p
317

318
        # Successors
319
        assert len(test2.getOperatorSuccessors()) == 0
320
        assert len(test1.getOperatorSuccessors()) == 1
321
        assert test1.getOperatorSuccessors()[0] == test2
322

323
        # Check all the nodes are valid (pybind ownership test)
324
        for node in [test1, test2, test3]:
325
            assert node.isOperator()
326
        for node in [x, y]:
327
            assert node.isTensor()
328

329
    def test_delete_node(self):
330
        nn = ng.NNModule()
331
        node = nn.dataFlow.createNode(ng.NeuralNetOperator("TestOp"))
332
        nn.dataFlow.deleteNode(node)
333
        assert len(nn.dataFlow.getMutableNodes()) == 0
334

335
    def test_replace_producer(self):
336
        net = core.Net("name")
337
        net.FC(["X", "W"], ["Y"])
338
        nn = ng.NNModule(net)
339
        fc = nn.controlFlow[0]
340
        test_op = nn.dataFlow.createNode(ng.NeuralNetOperator("TestOp"))
341
        nn.replaceProducer(fc.outputs[0], test_op)
342
        nn.deleteNode(fc)
343
        assert len(nn.controlFlow) == 1
344
        assert nn.controlFlow[0].name == "TestOp"
345

346
    def test_replace_all_uses_with(self):
347
        net = core.Net("name")
348
        net.FC(["X", "W"], ["Y"])
349
        net.FC(["X", "W2"], ["Y2"])
350
        nn = ng.NNModule(net)
351
        fc = nn.controlFlow[0]
352
        test_tensor = nn.dataFlow.createNode(ng.NeuralNetData("T"))
353
        nn.replaceAllUsesWith(fc.inputs[0], test_tensor)
354

355
        for op in nn.controlFlow:
356
            assert op.inputs[0].name == "T"
357

358
    def test_replace_as_consumer(self):
359
        net = core.Net("name")
360
        net.FC(["X", "W"], ["Y"])
361
        nn = ng.NNModule(net)
362
        fc = nn.controlFlow[0]
363
        test_op = nn.dataFlow.createNode(ng.NeuralNetOperator("TestOp"))
364
        nn.replaceAsConsumer(fc, test_op)
365
        nn.deleteNode(fc)
366
        assert len(nn.controlFlow) == 1
367
        assert nn.controlFlow[0].name == "TestOp"
368
        assert nn.controlFlow[0].inputs[0].name == "X"
369
        assert nn.controlFlow[0].inputs[1].name == "W"
370

371
    def test_annotation_basic(self):
372
        annot = ng.Annotation()
373
        annot.setDevice("woot")
374
        assert annot.getDevice() == "woot"
375
        annot.setDeviceType(7)
376
        assert annot.getDeviceType() == 7
377

378
    def test_annotation_from_graph(self):
379
        nn = ng.NNModule()
380
        node = nn.dataFlow.createNode(ng.NeuralNetOperator("TestOp"))
381
        annot = node.getAnnotation()
382
        annot.setDeviceType(7)
383
        node.setAnnotation(annot)
384
        new_annot = node.getAnnotation()
385
        assert new_annot.getDeviceType() == 7
386

387
    def test_annotation_operator_def(self):
388
        nn = ng.NNModule()
389
        opdef = core.CreateOperator("Conv", [], [], engine="SENTINEL")
390
        node = nn.dataFlow.createNode(opdef)
391
        assert node.annotation.operator_def.engine == "SENTINEL"
392
        opdef = core.CreateOperator("Conv", [], [], engine="NEW_SENTINEL")
393
        node.annotation.operator_def = opdef
394
        netdef = nn.convertToCaffe2Proto()
395
        assert len(netdef.op) == 1
396
        assert netdef.op[0].engine == "NEW_SENTINEL"
397

398
    def test_annotation_device_option(self):
399
        nn = ng.NNModule()
400
        node = nn.dataFlow.createNode(ng.NeuralNetOperator("TestOp"))
401
        d = caffe2_pb2.DeviceOption()
402
        d.node_name = "test"
403
        node.annotation.device_option = d
404
        # access in a different way
405
        d_2 = nn.controlFlow[0].annotation.device_option
406
        assert d == d_2
407

408
    def test_has_device_option(self):
409
        nn = ng.NNModule()
410
        node = nn.dataFlow.createNode(ng.NeuralNetOperator("TestOp"))
411
        assert not node.annotation.hasDeviceOption()
412
        d = caffe2_pb2.DeviceOption()
413
        node.annotation.device_option = d
414
        assert node.annotation.hasDeviceOption()
415

416
    def test_distributed_annotations(self):
417
        nn = ng.NNModule()
418
        key = nn.dataFlow.createNode(ng.NeuralNetData("key"))
419
        length = nn.dataFlow.createNode(ng.NeuralNetData("length"))
420
        node = nn.dataFlow.createNode(ng.NeuralNetOperator("TestOp"))
421

422
        annot = ng.Annotation()
423
        annot.setKeyNode(key)
424
        annot.setLengthNode(length)
425
        annot.setComponentLevels(["", "test", "woot"])
426

427
        node.setAnnotation(annot)
428

429
        new_annot = node.getAnnotation()
430
        #assert new_annot.getLengthNode() == length
431
        assert new_annot.getKeyNode() == key
432
        assert len(new_annot.getComponentLevels()) == 3
433
        assert new_annot.getComponentLevels()[0] == ""
434
        assert new_annot.getComponentLevels()[2] == "woot"
435

436
    def test_distributed_device_map(self):
437
        net = core.Net("name")
438
        net.FC(["X", "W"], ["Y"])
439
        d = caffe2_pb2.DeviceOption()
440
        nn = ng.NNModule(net, {"X": d, "W": d})
441

442
        with self.assertRaises(Exception):
443
            nn = ng.NNModule(net, {"X": d, "Fake": d})
444