def test_replace_as_consumer(self):
net = core.Net("name")
net.FC(["X", "W"], ["Y"])
nn = ng.NNModule(net)
fc = nn.controlFlow[0]
test_op = nn.dataFlow.createNode(ng.NeuralNetOperator("TestOp"))
nn.replaceAsConsumer(fc, test_op)
nn.deleteNode(fc)
assert len(nn.controlFlow) == 1
assert nn.controlFlow[0].name == "TestOp"
assert nn.controlFlow[0].inputs[0].name == "X"
assert nn.controlFlow[0].inputs[1].name == "W"
def test_match_graph_node(self):
mg = ng.NNMatchGraph()
mg.createNode(ng.NeuralNetOperator("test"))
nn = ng.NNModule()
test = nn.dataFlow.createNode(ng.NeuralNetOperator("test"))
x = nn.dataFlow.createNode(ng.NeuralNetData("X"))
nn.dataFlow.createEdge(x, test)
count = 0
for match in nn.match(mg):
assert len(match) == 1
count += 1
assert count == 1
def test_traversal(self):
net = core.Net("test")
net.FC(["X", "W"], ["Y"])
net.Relu(["Y"], ["Z"])
nn = ng.NNModule(net)
fc = nn.controlFlow[0]
relu = nn.controlFlow[1]
assert fc.inputs[0].name == "X"
assert fc.inputs[1].name == "W"
assert relu.outputs[0].name == "Z"
assert relu.inputs[0].name == "Y"
assert relu.inputs[0].producer.name == "FC"
assert fc.outputs[0].consumers[0].name == "Relu"
def test_createUniqueDataNode(self):
net = core.Net("name")
nn = ng.NNModule(net)
n1 = nn.createUniqueDataNode("a")
self.assertEqual(n1.name[0], "a")
n2 = nn.dataFlow.createNode(ng.Operator("test1"))
nn.createEdge(n1, n2)
n3 = nn.createUniqueDataNode("a")
nn.createEdge(n2, n3)
self.assertEqual(n3.name[0], "a")
self.assertNotEqual(n1.name, n3.name)
n1 = nn.createUniqueDataNode("b")
n2 = nn.createUniqueDataNode("b")
self.assertNotEqual(n1.name, n2.name)
def test_match_graph_node_strict(self):
mg = ng.NNMatchGraph()
mg.createNode(ng.NeuralNetOperator("test"), strict=True)
nn = ng.NNModule()
test = nn.dataFlow.createNode(ng.NeuralNetOperator("test"))
x = nn.dataFlow.createNode(ng.NeuralNetData("X"))
nn.dataFlow.createEdge(test, x)
count = 0
for match in nn.match(mg):
assert len(match) == 1
count += 1
with self.assertRaises(Exception):
assert count == 1
def test_match_graph_node(self):
mg = ng.NNMatchGraph()
mg.createNode(ng.NeuralNetOperator("test"))
nn = ng.NNModule()
test = nn.dataFlow.createNode(ng.NeuralNetOperator("test"))
x = nn.dataFlow.createNode(ng.NeuralNetData("X"))
nn.dataFlow.createEdge(x, test)
count = 0
for match in nn.match(mg):
assert len(match) == 1
count += 1
# Dot generation of subgraph
assert(str(match).startswith("digraph G"))
assert count == 1
def test_edges_complex(self, size):
random.seed(1337)
nn = ng.NNModule()
dfg = nn.dataFlow
data = []
ops = []
for _ in range(size):
data.append(dfg.createNode(ng.NeuralNetData("X")))
for i in range(size):
ops.append(dfg.createNode(ng.NeuralNetOperator("Op" + str(i))))
for i in range(size):
for j in range(size):
if bool(random.getrandbits(1)):
dfg.createEdge(data[i], ops[j])
def test_core_net_nn_accessors(self):
net = core.Net("name")
net.FC(["X", "W"], ["Y"])
net.Relu(["Y"], ["Z"])
nn = ng.NNModule(net)
tensors = set()
for t in nn.tensors:
tensors.add(t.name)
assert tensors == set(["X", "W", "Y", "Z"])
ops = set()
for op in nn.operators:
ops.add(op.name)
assert ops == set(["FC", "Relu"])
nodes = set()
for node in nn.nodes:
nodes.add(node.name)
assert nodes == (ops | tensors)
def test_simple_rewire(self):
net = core.Net("name")
# Rewire this so that we get
# c = Add(a, d)
# e = Mul(c, b)
#
# if a = 1, b = 2, d = 3
# we get 8: (1 + 3) * 2
# as opposed to 7: 1 + (3 * 2)
net.Mul(["a", "b"], ["c"])
net.Add(["c", "d"], ["e"])
nn = ng.NNModule(net)
mul = nn.controlFlow[0]
add = nn.controlFlow[1]
a = mul.inputs[0]
b = mul.inputs[1]
c = mul.outputs[0]
d = add.inputs[1]
e = add.outputs[0]
nn.deleteEdge(a, mul)
nn.deleteEdge(b, mul)
nn.deleteEdge(mul, c)
nn.deleteEdge(c, add)
nn.deleteEdge(d, add)
nn.deleteEdge(add, e)
nn.createEdge(a, add)
nn.createEdge(d, add)
nn.createEdge(add, c)
nn.createEdge(c, mul)
nn.createEdge(b, mul)
nn.createEdge(mul, e)
# Test it out
new_netdef = nn.convertToCaffe2Proto()
workspace.ResetWorkspace()
workspace.FeedBlob("a", np.array([1, 1, 1]))
workspace.FeedBlob("b", np.array([2, 2, 2]))
workspace.FeedBlob("d", np.array([3, 3, 3]))
workspace.RunNetOnce(new_netdef)
out = workspace.FetchBlob("e")
expected_out = np.array([8, 8, 8])
np.testing.assert_almost_equal(out, expected_out)
def test_simple_replace(self):
net = core.Net("name")
net.FC(["X", "W"], ["Y"])
nn = ng.NNModule(net)
fc = nn.controlFlow[0]
add = nn.createNode(core.CreateOperator("Add", ["X"], ["Y"], engine="CUDNN"))
nn.replaceNode(fc, add)
nn.deleteNode(fc)
# Test it out
new_netdef = nn.convertToCaffe2Proto()
workspace.ResetWorkspace()
workspace.FeedBlob("X", np.array([1, 2, 3]))
workspace.FeedBlob("W", np.array([1, 2, 3]))
workspace.RunNetOnce(new_netdef)
out = workspace.FetchBlob("Y")
expected_out = np.array([2, 4, 6])
np.testing.assert_almost_equal(out, expected_out)
def test_delete_subgraph(self):
mg = ng.NNMatchGraph()
test2m = mg.createNode(ng.NeuralNetOperator("test2"), strict=True)
xm = mg.createNode(ng.NeuralNetData("X"), strict=True)
testm = mg.createNode(ng.NeuralNetOperator("test"))
mg.createEdge(test2m, xm)
mg.createEdge(xm, testm)
nn = ng.NNModule()
test2 = nn.dataFlow.createNode(ng.NeuralNetOperator("test2"))
x = nn.dataFlow.createNode(ng.NeuralNetData("X"))
test = nn.dataFlow.createNode(ng.NeuralNetOperator("test"))
nn.dataFlow.createEdge(test2, x)
nn.dataFlow.createEdge(x, test)
for m in nn.match(mg):
match = m
nn.deleteSubgraph(match)
assert len(nn.controlFlow) == 0
def test_distributed_annotations(self):
nn = ng.NNModule()
key = nn.dataFlow.createNode(ng.NeuralNetData("key"))
length = nn.dataFlow.createNode(ng.NeuralNetData("length"))
node = nn.dataFlow.createNode(ng.NeuralNetOperator("TestOp"))
annot = ng.Annotation()
annot.setKeyNode(key)
annot.setLengthNode(length)
annot.setComponentLevels(["", "test", "woot"])
node.setAnnotation(annot)
new_annot = node.getAnnotation()
#assert new_annot.getLengthNode() == length
assert new_annot.getKeyNode() == key
assert len(new_annot.getComponentLevels()) == 3
assert new_annot.getComponentLevels()[0] == ""
assert new_annot.getComponentLevels()[2] == "woot"
def test_replace_subraph(self):
mg = ng.NNMatchGraph()
test2m = mg.createNode(ng.NeuralNetOperator("test2"), strict=True)
xm = mg.createNode(ng.NeuralNetData("X"), strict=True)
testm = mg.createNode(ng.NeuralNetOperator("test"))
mg.createEdge(test2m, xm)
mg.createEdge(xm, testm)
nn = ng.NNModule()
test2 = nn.dataFlow.createNode(ng.NeuralNetOperator("test2"))
x = nn.dataFlow.createNode(ng.NeuralNetData("X"))
test = nn.dataFlow.createNode(ng.NeuralNetOperator("test"))
nn.dataFlow.createEdge(test2, x)
nn.dataFlow.createEdge(x, test)
for m in nn.match(mg):
match = m
new_op = nn.dataFlow.createNode(ng.NeuralNetOperator("new_op"))
nn.replaceSubgraph(match, new_op, [], [])
assert len(nn.controlFlow) == 1
assert nn.controlFlow[0].name == "new_op"
def test_match_graph(self):
mg = ng.NNMatchGraph()
test2m = mg.createNode(ng.NeuralNetOperator("test2"), strict=True)
xm = mg.createNode(ng.NeuralNetData("X"), strict=True)
testm = mg.createNode(ng.NeuralNetOperator("test"))
mg.createEdge(test2m, xm)
mg.createEdge(xm, testm)
nn = ng.NNModule()
test2 = nn.dataFlow.createNode(ng.NeuralNetOperator("test2"))
x = nn.dataFlow.createNode(ng.NeuralNetData("X"))
test = nn.dataFlow.createNode(ng.NeuralNetOperator("test"))
nn.dataFlow.createEdge(test2, x)
nn.dataFlow.createEdge(x, test)
count = 0
for match in nn.match(mg):
print(len(match))
assert len(match) == 3
count += 1
assert count == 1
def test_convertToProto(self):
net = core.Net("name")
net.FC(["X", "W"], ["Y"])
nn = ng.NNModule(net)
new_netdef = nn.convertToCaffe2Proto()
print(new_netdef)
print(net.Proto())
assert len(new_netdef.op) == len(net.Proto().op)
for i in range(len(new_netdef.op)):
op = net.Proto().op[i]
new_op = new_netdef.op[i]
assert op.type == new_op.type
assert len(op.input) == len(new_op.input)
assert len(op.output) == len(new_op.output)
for a, b in zip(op.input, new_op.input):
assert a == b
for a, b in zip(op.output, new_op.output):
assert a == b
for a, b in zip(new_netdef.external_input, net.Proto().external_input):
assert a == b
for a, b in zip(new_netdef.external_output, net.Proto().external_output):
assert a == b
def test_edges_simple(self):
nn = ng.NNModule()
dfg = nn.dataFlow
x = dfg.createNode(ng.NeuralNetData("X"))
w = dfg.createNode(ng.NeuralNetData("W"))
op = dfg.createNode(ng.NeuralNetOperator("Op"))
with self.assertRaises(Exception):
dfg.createEdge(x, w)
dfg.createEdge(op, w)
dfg.createEdge(x, op)
# Dot generation
assert(str(dfg).startswith("digraph G"))
# subgraph
sg = ng.NNSubgraph()
sg.addNode(x)
sg.addNode(op)
sg.induceEdges()
assert len(sg) == 2
# subgraph dot generation
assert(str(sg).startswith("digraph G"))
def test_invalid_node(self):
nn = ng.NNModule()
dfg = nn.dataFlow
with self.assertRaises(Exception):
dfg.createNode(7)
def test_simple(self):
nn = ng.NNModule()
dfg = nn.dataFlow
dfg.createNode(ng.NeuralNetData("X"))
dfg.createNode(ng.NeuralNetOperator("FC"))
assert len(nn.dataFlow.getMutableNodes()) == 2
def test_delete_node(self):
nn = ng.NNModule()
node = nn.dataFlow.createNode(ng.NeuralNetOperator("TestOp"))
nn.dataFlow.deleteNode(node)
assert len(nn.dataFlow.getMutableNodes()) == 0
