In PyTorch nn.Linear has to take 2D tensor, this often leads to reshape a 4D tensor to 2D by calling .view(). However this (dynamic) reshaping doesn't work well with ONNX and Int8 tools, and cause using extra ops (eg. ExpandDims) that might not be available on mobile.
(predict_net, params)
| 880 | |
| 881 | |
| 882 | def remove_reshape_for_fc(predict_net, params): |
| 883 | """ |
| 884 | In PyTorch nn.Linear has to take 2D tensor, this often leads to reshape |
| 885 | a 4D tensor to 2D by calling .view(). However this (dynamic) reshaping |
| 886 | doesn't work well with ONNX and Int8 tools, and cause using extra |
| 887 | ops (eg. ExpandDims) that might not be available on mobile. |
| 888 | Luckily Caffe2 supports 4D tensor for FC, so we can remove those reshape |
| 889 | after exporting ONNX model. |
| 890 | """ |
| 891 | from caffe2.python import core |
| 892 | |
| 893 | # find all reshape sub-graph that can be removed, which is now all Reshape |
| 894 | # sub-graph whose output is only consumed by FC. |
| 895 | # TODO: to make it safer, we may need the actually value to better determine |
| 896 | # if a Reshape before FC is removable. |
| 897 | reshape_sub_graphs = identify_reshape_sub_graph(predict_net) |
| 898 | sub_graphs_to_remove = [] |
| 899 | for reshape_sub_graph in reshape_sub_graphs: |
| 900 | reshape_op_id = reshape_sub_graph[-1] |
| 901 | assert predict_net.op[reshape_op_id].type == "Reshape" |
| 902 | ssa, _ = core.get_ssa(predict_net) |
| 903 | reshape_output = ssa[reshape_op_id][1][0] |
| 904 | consumers = [i for i in range(len(ssa)) if reshape_output in ssa[i][0]] |
| 905 | if all(predict_net.op[consumer].type == "FC" for consumer in consumers): |
| 906 | # safety check if the sub-graph is isolated, for this reshape sub-graph, |
| 907 | # it means it has one non-param external input and one external output. |
| 908 | ext_inputs, ext_outputs = get_sub_graph_external_input_output( |
| 909 | predict_net, reshape_sub_graph |
| 910 | ) |
| 911 | non_params_ext_inputs = [inp for inp in ext_inputs if inp[1] != 0] |
| 912 | if len(non_params_ext_inputs) == 1 and len(ext_outputs) == 1: |
| 913 | sub_graphs_to_remove.append(reshape_sub_graph) |
| 914 | |
| 915 | # perform removing subgraph by: |
| 916 | # 1: rename the Reshape's output to its input, then the graph can be |
| 917 | # seen as in-place itentify, meaning whose external input/output are the same. |
| 918 | # 2: simply remove those ops. |
| 919 | remove_op_ids = [] |
| 920 | params_to_remove = [] |
| 921 | for sub_graph in sub_graphs_to_remove: |
| 922 | logger.info( |
| 923 | "Remove Reshape sub-graph:\n{}".format( |
| 924 | "".join(["(#{:>4})\n{}".format(i, predict_net.op[i]) for i in sub_graph]) |
| 925 | ) |
| 926 | ) |
| 927 | reshape_op_id = sub_graph[-1] |
| 928 | new_reshap_output = predict_net.op[reshape_op_id].input[0] |
| 929 | rename_op_output(predict_net, reshape_op_id, 0, new_reshap_output) |
| 930 | ext_inputs, ext_outputs = get_sub_graph_external_input_output(predict_net, sub_graph) |
| 931 | non_params_ext_inputs = [inp for inp in ext_inputs if inp[1] != 0] |
| 932 | params_ext_inputs = [inp for inp in ext_inputs if inp[1] == 0] |
| 933 | assert len(non_params_ext_inputs) == 1 and len(ext_outputs) == 1 |
| 934 | assert ext_outputs[0][0] == non_params_ext_inputs[0][0] |
| 935 | assert ext_outputs[0][1] == non_params_ext_inputs[0][1] + 1 |
| 936 | remove_op_ids.extend(sub_graph) |
| 937 | params_to_remove.extend(params_ext_inputs) |
| 938 | |
| 939 | predict_net = copy.deepcopy(predict_net) |
no test coverage detected