def replace_addv2(graph): """Replace all 'AddV2' in the graph with 'Add'. NOTE: 'AddV2' is not supported by UFF parser. """ for node in graph.find_nodes_by_op('AddV2'): gs.update_node(node, op='Add') return graph
def replace_fusedbnv3(graph): """Replace all 'FusedBatchNormV3' in the graph with 'FusedBatchNorm'. NOTE: 'FusedBatchNormV3' is not supported by UFF parser. https://devtalk.nvidia.com/default/topic/1066445/tensorrt/tensorrt-6-0-1-tensorflow-1-14-no-conversion-function-registered-for-layer-fusedbatchnormv3-yet/post/5403567/#5403567 """ for node in graph.find_nodes_by_op('FusedBatchNormV3'): gs.update_node(node, op='FusedBatchNorm') return graph
def replace_addv2(graph): """Replace all 'AddV2' in the graph with 'Add'. 'AddV2' is not supported by UFF parser. Reference: 1. https://github.com/jkjung-avt/tensorrt_demos/issues/113#issuecomment-629900809 """ for node in graph.find_nodes_by_op('AddV2'): gs.update_node(node, op='Add') return graph
def replace_fusedbnv3(graph): """Replace all 'FusedBatchNormV3' in the graph with 'FusedBatchNorm'. 'FusedBatchNormV3' is not supported by UFF parser. Reference: 1. https://devtalk.nvidia.com/default/topic/1066445/tensorrt/tensorrt-6-0-1-tensorflow-1-14-no-conversion-function-registered-for-layer-fusedbatchnormv3-yet/post/5403567/#5403567 2. https://github.com/jkjung-avt/tensorrt_demos/issues/76#issuecomment-607879831 """ for node in graph.find_nodes_by_op('FusedBatchNormV3'): gs.update_node(node, op='FusedBatchNorm') return graph
def add_plugin(graph): all_assert_nodes = graph.find_nodes_by_op("Assert") graph.remove(all_assert_nodes, remove_exclusive_dependencies=True) all_identity_nodes = graph.find_nodes_by_op("Identity") graph.forward_inputs(all_identity_nodes) Input = gs.create_plugin_node(name="Input", op="Placeholder", shape=[1, 3, 300, 300]) PriorBox = gs.create_plugin_node( name="GridAnchor", op="GridAnchor_TRT", minSize=0.2, maxSize=0.95, aspectRatios=[1.0, 2.0, 0.5, 3.0, 0.33], variance=[0.1, 0.1, 0.2, 0.2], featureMapShapes=[19, 10, 5, 3, 2, 1], # Resolution 300 #featureMapShapes=[29, 15, 8, 4, 2, 1], # Resolution 450 numLayers=6) NMS = gs.create_plugin_node(name="NMS", op="NMS_TRT", shareLocation=1, varianceEncodedInTarget=0, backgroundLabelId=0, confidenceThreshold=1e-8, nmsThreshold=0.6, topK=100, keepTopK=100, numClasses=2, inputOrder=[0, 2, 1], confSigmoid=1, isNormalized=1) concat_priorbox = gs.create_node("concat_priorbox", op="ConcatV2", axis=2) concat_box_loc = gs.create_plugin_node("concat_box_loc", op="FlattenConcat_TRT", axis=1, ignoreBatch=0) concat_box_conf = gs.create_plugin_node("concat_box_conf", op="FlattenConcat_TRT", axis=1, ignoreBatch=0) namespace_plugin_map = { "MultipleGridAnchorGenerator": PriorBox, "Postprocessor": NMS, "Preprocessor": Input, "Cast": Input, "image_tensor": Input, "Concatenate": concat_priorbox, "concat": concat_box_loc, "concat_1": concat_box_conf } for node in graph.find_nodes_by_op('FusedBatchNormV3'): gs.update_node(node, op='FusedBatchNorm') graph.collapse_namespaces(namespace_plugin_map) graph.remove(graph.graph_outputs, remove_exclusive_dependencies=False) graph.find_nodes_by_op("NMS_TRT")[0].input.remove("Input") # Create a constant Tensor and set it as input for GridAnchor_TRT data = np.array([1, 1], dtype=np.float32) anchor_input = gs.create_node("AnchorInput", "Const", value=data) graph.append(anchor_input) graph.find_nodes_by_op("GridAnchor_TRT")[0].input.insert(0, "AnchorInput") return graph