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.
(nn)
| 7 | |
| 8 | |
| 9 | def transpose_network(nn): |
| 10 | """ |
| 11 | Convert all Convolutions operators which are in the NCHW order |
| 12 | to NHWC order and also transform their inputs and outputs so that the |
| 13 | rest of the graph is not affected. |
| 14 | """ |
| 15 | # track the incoming tensors into NHWC2NCHW operators |
| 16 | incoming = {} # output tensor -> input tensor |
| 17 | # track outgoing tensors from NCHW2NHWC operators |
| 18 | outgoing = defaultdict(lambda: []) # input tensor -> list of operators |
| 19 | dfg = nn.dataFlow |
| 20 | orig_nodes = [x for x in nn.nodes] |
| 21 | for node in orig_nodes: |
| 22 | if node.isOperator() and node.name == "Conv": |
| 23 | arg_dict = utils.ArgsToDict(node.annotation.operator_def.arg) |
| 24 | # a missing "order" argument implies default NCHW order |
| 25 | if "order" in arg_dict and arg_dict["order"] != "NCHW": |
| 26 | continue |
| 27 | inputs = [x for x in node.inputs] |
| 28 | assert len(inputs) >= 2, "Conv operator should have two inputs" |
| 29 | outputs = [x for x in node.outputs] |
| 30 | assert len(outputs) >= 1, "Conv operator should have an output" |
| 31 | for inp in inputs: |
| 32 | nn.deleteEdge(inp, node) |
| 33 | for outp in outputs: |
| 34 | nn.deleteEdge(node, outp) |
| 35 | # only the first two inputs of the Convolution the data and the |
| 36 | # weights need to be transformed |
| 37 | for idx in range(2): |
| 38 | new_inp = nn.createUniqueDataNode(inputs[idx].name) |
| 39 | transp = dfg.createNode(ng.NeuralNetOperator("NCHW2NHWC")) |
| 40 | nn.createEdge(inputs[idx], transp) |
| 41 | nn.createEdge(transp, new_inp) |
| 42 | outgoing[inputs[idx]].append(transp) |
| 43 | inputs[idx] = new_inp |
| 44 | for idx in range(len(outputs)): |
| 45 | new_outp = nn.createUniqueDataNode(outputs[idx].name) |
| 46 | transp = dfg.createNode(ng.NeuralNetOperator("NHWC2NCHW")) |
| 47 | nn.createEdge(transp, outputs[idx]) |
| 48 | nn.createEdge(new_outp, transp) |
| 49 | incoming[outputs[idx]] = new_outp |
| 50 | outputs[idx] = new_outp |
| 51 | # create a new Convolution with identical arguments as the original |
| 52 | # one except for the order |
| 53 | arg_dict["order"] = "NHWC" |
| 54 | new_node = nn.createNode(core.CreateOperator("Conv", [], [], |
| 55 | **arg_dict)) |
| 56 | for inp in inputs: |
| 57 | nn.createEdge(inp, new_node) |
| 58 | for outp in outputs: |
| 59 | nn.createEdge(new_node, outp) |
| 60 | |
| 61 | nn.deleteNode(node) |
| 62 | |
| 63 | # finally, we will compress |
| 64 | # case 1: |
| 65 | # X -> NHWC2NCHW -> Y -> NCHW2NHWC -> Z1 ; Y -> NCHW2NHWC -> Z2 |
| 66 | # to: |
searching dependent graphs…