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
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 transpose_network(nn): """ Convert all Convolutions operators which are in the NCHW order to NHWC order and also transform their inputs and outputs so that the rest of the graph is not affected. """ # track the incoming tensors into NHWC2NCHW operators incoming = {} # output tensor -> input tensor # track outgoing tensors from NCHW2NHWC operators outgoing = defaultdict(lambda: []) # input tensor -> list of operators dfg = nn.dataFlow orig_nodes = [x for x in nn.nodes] for node in orig_nodes: if node.isOperator() and node.name == "Conv": arg_dict = utils.ArgsToDict(node.annotation.operator_def.arg) # a missing "order" argument implies default NCHW order if "order" in arg_dict and arg_dict["order"] != "NCHW": continue inputs = [x for x in node.inputs] assert len(inputs) >= 2, "Conv operator should have two inputs" outputs = [x for x in node.outputs] assert len(outputs) >= 1, "Conv operator should have an output" for inp in inputs: nn.deleteEdge(inp, node) for outp in outputs: nn.deleteEdge(node, outp) # only the first two inputs of the Convolution the data and the # weights need to be transformed for idx in range(2): new_inp = nn.createUniqueDataNode(inputs[idx].name) transp = dfg.createNode(ng.NeuralNetOperator("NCHW2NHWC")) nn.createEdge(inputs[idx], transp) nn.createEdge(transp, new_inp) outgoing[inputs[idx]].append(transp) inputs[idx] = new_inp for idx in range(len(outputs)): new_outp = nn.createUniqueDataNode(outputs[idx].name) transp = dfg.createNode(ng.NeuralNetOperator("NHWC2NCHW")) nn.createEdge(transp, outputs[idx]) nn.createEdge(new_outp, transp) incoming[outputs[idx]] = new_outp outputs[idx] = new_outp # create a new Convolution with identical arguments as the original # one except for the order arg_dict["order"] = "NHWC" new_node = nn.createNode(core.CreateOperator("Conv", [], [], **arg_dict)) for inp in inputs: nn.createEdge(inp, new_node) for outp in outputs: nn.createEdge(new_node, outp) nn.deleteNode(node) # finally, we will compress # case 1: # X -> NHWC2NCHW -> Y -> NCHW2NHWC -> Z1 ; Y -> NCHW2NHWC -> Z2 # to: # X -> NHWC2NCHW -> Y and replace Z1 with X and replace Z2 with X # And case 2: # Y -> NCHW2NHWC -> Z1 ; Y -> NCHW2NHWC -> Z2 # to: # Y -> NCHW2NHWC -> Z1 and replace Z2 with Z1 # orig_tensor is one of the tensors in the original graph in NCHW order for orig_tensor in outgoing: # new_tensor is identical to orig_tensor except the order is NHWC if orig_tensor in incoming: # case 1 (see above) new_tensor = incoming[orig_tensor] else: # case 2 (see above) out_ops = outgoing[orig_tensor] new_tensor = out_ops[0].outputs[0] outgoing[orig_tensor] = out_ops[1:] for opnode in outgoing[orig_tensor]: # there should only be one output, so this iteration is overkill for out in opnode.outputs: nn.replaceAllUsesWith(out, new_tensor) nn.deleteNode(out) nn.deleteNode(opnode)