| 30 | TRT_MODEL_PATH = os.path.join(MODEL_DIR, "bidaf-9-trt.onnx") |
| 31 | |
| 32 | def _do_graph_surgery(raw_model_path, trt_model_path): |
| 33 | graph = gs.import_onnx(onnx.load(raw_model_path)) |
| 34 | |
| 35 | # Replace unsupported Hardmax with our CustomHardmax op |
| 36 | for node in graph.nodes: |
| 37 | if node.op == 'Hardmax': |
| 38 | node.op = 'CustomHardmax' |
| 39 | hardmax_node = node |
| 40 | |
| 41 | # The original onnx model also uses another unsupported op called "Compress". |
| 42 | # "Compress" returns values from the first tensor for all indices which evaluate to |
| 43 | # True in the second tensor. In our case the second Tensor is the output of Hardmax, |
| 44 | # so exactly one index will evaluate to true because the value at it will be 1, and |
| 45 | # all other values will be 0. We can achieve the same result as "Compress" by taking the |
| 46 | # dot product of our value tensor and the Hardmax output. |
| 47 | # |
| 48 | # So, we will replace the subgraph Compress(Transpose_29, Cast(Reshape(Hardmax))) |
| 49 | # with the subgraph Einsum(Transpose_29, Hardmax) where the equation in Einsum takes the dot product. |
| 50 | node_by_name = {node.name : node for node in graph.nodes} |
| 51 | transpose_node = node_by_name['Transpose_29'] |
| 52 | compress_node = node_by_name['Compress_31'] |
| 53 | |
| 54 | einsum_node = gs.Node( |
| 55 | 'Einsum', |
| 56 | 'Dot_of_Hardmax_and_Transpose', |
| 57 | attrs={'equation': 'ij,ij->i'}, # "Dot product" of 2d tensors |
| 58 | inputs=[hardmax_node.outputs[0], transpose_node.outputs[0]], |
| 59 | outputs=[compress_node.outputs[0]] |
| 60 | ) |
| 61 | graph.nodes.append(einsum_node) |
| 62 | |
| 63 | # Separate the old subgraph which will be deleted with graph.cleanup() |
| 64 | hardmax_node.o().inputs.clear() |
| 65 | transpose_node.o().inputs.clear() |
| 66 | compress_node.outputs.clear() |
| 67 | |
| 68 | # Also remove the CategoryMapper nodes which convert strings to integers as the first step in the model. |
| 69 | # We need to convert the following structure: |
| 70 | # |
| 71 | # Input as Converted to |
| 72 | # String tokens Integer tokens |
| 73 | # ---------------->[CategoryMapper]------------------>[Rest of Model] |
| 74 | # |
| 75 | # into the following: |
| 76 | # |
| 77 | # Input as |
| 78 | # Integer tokens |
| 79 | # ------------------>[Rest of Model] |
| 80 | # |
| 81 | # Later we will feed the model the integer tokens directly. |
| 82 | # Note: list conversion is necessary because we modify graph.nodes in the for loop. |
| 83 | category_mapper_nodes = [node for node in graph.nodes if node.op == 'CategoryMapper'] |
| 84 | for node in category_mapper_nodes: |
| 85 | # Remove CategoryMapper node from onnx graph |
| 86 | graph.nodes.remove(node) |
| 87 | |
| 88 | # Also remove references its inputs in the graph's inputs |
| 89 | for input_tensor in node.inputs: |