def convert_unsupported_nodes_to_plugins(dynamic_graph):
Input = gs.create_node("Input",
op="Placeholder",
dtype=tf.float32,
shape=[1, 3, 300, 300])
PriorBox = gs.create_plugin_node(name="GridAnchor",
op="GridAnchor_TRT",
numLayers=6,
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])
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=91,
inputOrder=[0, 2, 1],
confSigmoid=1,
isNormalized=1)
concat_priorbox = gs.create_node(name="concat_priorbox",
op="ConcatV2",
dtype=tf.float32,
axis=2)
concat_box_loc = gs.create_plugin_node("concat_box_loc",
op="FlattenConcat_TRT",
dtype=tf.float32,
axis=1,
ignoreBatch=0)
concat_box_conf = gs.create_plugin_node("concat_box_conf",
op="FlattenConcat_TRT",
dtype=tf.float32,
axis=1,
ignoreBatch=0)
namespace_plugin_map = {
"MultipleGridAnchorGenerator": PriorBox,
"Postprocessor": NMS,
"Preprocessor": Input,
"ToFloat": Input,
"image_tensor": Input,
"MultipleGridAnchorGenerator/Concatenate": concat_priorbox,
"MultipleGridAnchorGenerator/Identity": concat_priorbox,
"concat": concat_box_loc,
"concat_1": concat_box_conf
}
# Now create a new graph by collapsing namespaces
dynamic_graph.collapse_namespaces(namespace_plugin_map)
# Remove the outputs, so we just have a single output node (NMS).
dynamic_graph.remove(dynamic_graph.graph_outputs,
remove_exclusive_dependencies=False)
return dynamic_graph
def add_anchor_input(graph):
"""Add the missing const input for the GridAnchor node.
Reference:
1. https://www.minds.ai/post/deploying-ssd-mobilenet-v2-on-the-nvidia-jetson-and-nano-platforms
"""
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
def parse_gridAnchor(graph):
"""
define a constant input tensor and set that as the input for the GridAnchor node
as UFF file does not provide an input element for the GridAnchor node
"""
data = np.array([1, 1], dtype=np.float32)
anchor_input = gs.create_node("AnchorInput", "Const", value=data, dtype=tf.float32)
graph.append(anchor_input)
graph.find_nodes_by_op("GridAnchor_TRT")[0].input.insert(0, "AnchorInput")
return graph
def preprocess(dynamic_graph):
# Now create a new graph by collapsing namespaces
dynamic_graph.collapse_namespaces(namespace_plugin_map)
# Remove the outputs, so we just have a single output node (NMS).
dynamic_graph.remove(dynamic_graph.graph_outputs,
remove_exclusive_dependencies=False)
# 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")
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],
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=91,
inputOrder=[1, 0, 2],
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",
)
concat_box_conf = gs.create_plugin_node(
"concat_box_conf",
op="FlattenConcat_TRT",
)
namespace_plugin_map = {
"MultipleGridAnchorGenerator": PriorBox,
"Postprocessor": NMS,
"Preprocessor": Input,
"ToFloat": Input,
"image_tensor": Input,
"Concatenate": concat_priorbox,
"concat": concat_box_loc,
"concat_1": concat_box_conf
}
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")
return graph
def ssd_unsupported_nodes_to_plugin_nodes(ssd_graph):
"""Makes ssd_graph TensorRT comparible using graphsurgeon.
This function takes ssd_graph, which contains graphsurgeon
DynamicGraph data structure. This structure describes frozen Tensorflow
graph, that can be modified using graphsurgeon (by deleting, adding,
replacing certain nodes). The graph is modified by removing
Tensorflow operations that are not supported by TensorRT's UffParser
and replacing them with custom layer plugin nodes.
Note: This specific implementation works only for
ssd_inception_v2_coco_2017_11_17 network.
Args:
ssd_graph (gs.DynamicGraph): graph to convert
Returns:
gs.DynamicGraph: UffParser compatible SSD graph
"""
# Create TRT plugin nodes to replace unsupported ops in Tensorflow graph
channels = ModelData.get_input_channels()
height = ModelData.get_input_height()
width = ModelData.get_input_width()
Input = gs.create_plugin_node(name="Input",
op="Placeholder",
dtype=tf.float32,
shape=[1, channels, height, width])
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],
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=91,
inputOrder=[0, 2, 1],
confSigmoid=1,
isNormalized=1)
concat_priorbox = gs.create_node("concat_priorbox",
op="ConcatV2",
dtype=tf.float32,
axis=2)
concat_box_loc = gs.create_plugin_node("concat_box_loc",
op="FlattenConcat_TRT",
dtype=tf.float32,
axis=1,
ignoreBatch=0)
concat_box_conf = gs.create_plugin_node("concat_box_conf",
op="FlattenConcat_TRT",
dtype=tf.float32,
axis=1,
ignoreBatch=0)
# Create a mapping of namespace names -> plugin nodes.
namespace_plugin_map = {
"MultipleGridAnchorGenerator": PriorBox,
"Postprocessor": NMS,
"Preprocessor": Input,
"ToFloat": Input,
"image_tensor": Input,
"MultipleGridAnchorGenerator/Concatenate": concat_priorbox,
"MultipleGridAnchorGenerator/Identity": concat_priorbox,
"concat": concat_box_loc,
"concat_1": concat_box_conf
}
# Create a new graph by collapsing namespaces
ssd_graph.collapse_namespaces(namespace_plugin_map)
# Remove the outputs, so we just have a single output node (NMS).
# If remove_exclusive_dependencies is True, the whole graph will be removed!
ssd_graph.remove(ssd_graph.graph_outputs,
remove_exclusive_dependencies=False)
return ssd_graph
def add_plugin_and_preprocess(graph, model, num_classes, neuralet_adaptive_model):
"""add_plugin
Reference:
1. https://github.com/AastaNV/TRT_object_detection/blob/master/config/model_ssd_mobilenet_v1_coco_2018_01_28.py
2. https://github.com/AastaNV/TRT_object_detection/blob/master/config/model_ssd_mobilenet_v2_coco_2018_03_29.py
3. https://devtalk.nvidia.com/default/topic/1050465/jetson-nano/how-to-write-config-py-for-converting-ssd-mobilenetv2-to-uff-format/post/5333033/#5333033
"""
num_classes += 1
min_size = 0.2
max_size = 0.95
input_order = [0, 2, 1]
if neuralet_adaptive_model == 0:
input_order = [1, 0, 2]
input_dims = (3, 300, 300)
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,) + input_dims,
dtype=tf.float32
)
PriorBox = gs.create_plugin_node(
name="MultipleGridAnchorGenerator",
op="GridAnchor_TRT",
minSize=min_size, # was 0.2
maxSize=max_size, # was 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],
numLayers=6,
dtype=tf.float32
)
NMS = gs.create_plugin_node(
name="NMS",
op="NMS_TRT",
shareLocation=1,
varianceEncodedInTarget=0,
backgroundLabelId=0,
confidenceThreshold=0.3, # was 1e-8
nmsThreshold=0.6,
topK=100,
keepTopK=100,
numClasses=num_classes, # was 91
inputOrder=input_order,
confSigmoid=1,
isNormalized=1,
dtype=tf.float32
)
concat_priorbox = gs.create_node(
"concat_priorbox",
op="ConcatV2",
axis=2,
dtype=tf.float32
)
if trt.__version__[0] >= '6':
concat_box_loc = gs.create_plugin_node(
"concat_box_loc",
op="FlattenConcat_TRT",
axis=1,
ignoreBatch=0,
dtype=tf.float32
)
concat_box_conf = gs.create_plugin_node(
"concat_box_conf",
op="FlattenConcat_TRT",
axis=1,
ignoreBatch=0,
dtype=tf.float32
)
else:
concat_box_loc = gs.create_plugin_node(
"concat_box_loc",
op="FlattenConcat_TRT",
dtype=tf.float32
)
concat_box_conf = gs.create_plugin_node(
"concat_box_conf",
op="FlattenConcat_TRT",
dtype=tf.float32
)
namespace_plugin_map = {
"MultipleGridAnchorGenerator": PriorBox,
"Postprocessor": NMS,
"Preprocessor": Input,
"Cast": Input,
"ToFloat": Input,
"image_tensor": Input,
"MultipleGridAnchorGenerator/Concatenate": concat_priorbox, # for 'ssd_mobilenet_v1_coco'
"Concatenate": concat_priorbox,
"concat": concat_box_loc,
"concat_1": concat_box_conf
}
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")
if ( "image_tensor:0" in graph.find_nodes_by_name("Input")[0].input ):
# for ssd_mobilenet_v1
graph.find_nodes_by_name("Input")[0].input.remove("image_tensor:0")
graph = replace_addv2(graph)
graph = replace_fusedbnv3(graph)
graph = parse_gridAnchor(graph)
return graph
import graphsurgeon as gs
import tensorflow as tf
Input = gs.create_node("conv2d_input",
op="Placeholder",
dtype=tf.float32,
shape=[-1, 1, 28, 28])
isrlu_activation = gs.create_node(name="activation_1",
op="ISRLU",
trt_plugin=True,
dtype=tf.float32)
namespace_plugin_map = {"activation_1": ISRLU}
def preprocess(dynamic_graph):
# Now create a new graph by collapsing namespaces
dynamic_graph.collapse_namespaces(namespace_plugin_map)
# Remove the outputs, so we just have a single output node
dynamic_graph.remove(dynamic_graph.graph_outputs,
remove_exclusive_dependencies=False)
def add_plugin(graph, model, spec):
"""add_plugin
Reference:
1. https://github.com/AastaNV/TRT_object_detection/blob/master/config/model_ssd_mobilenet_v1_coco_2018_01_28.py
2. https://github.com/AastaNV/TRT_object_detection/blob/master/config/model_ssd_mobilenet_v2_coco_2018_03_29.py
3. https://devtalk.nvidia.com/default/topic/1050465/jetson-nano/how-to-write-config-py-for-converting-ssd-mobilenetv2-to-uff-format/post/5333033/#5333033
"""
numClasses = spec['num_classes']
minSize = spec['min_size']
maxSize = spec['max_size']
inputOrder = spec['input_order']
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,) + INPUT_DIMS
)
PriorBox = gs.create_plugin_node(
name="MultipleGridAnchorGenerator",
op="GridAnchor_TRT",
minSize=minSize, # was 0.2
maxSize=maxSize, # was 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],
numLayers=6
)
NMS = gs.create_plugin_node(
name="NMS",
op="NMS_TRT",
shareLocation=1,
varianceEncodedInTarget=0,
backgroundLabelId=0,
confidenceThreshold=0.3, # was 1e-8
nmsThreshold=0.6,
topK=100,
keepTopK=100,
numClasses=numClasses, # was 91
inputOrder=inputOrder,
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",
)
concat_box_conf = gs.create_plugin_node(
"concat_box_conf",
op="FlattenConcat_TRT",
)
namespace_plugin_map = {
"MultipleGridAnchorGenerator": PriorBox,
"Postprocessor": NMS,
"Preprocessor": Input,
"ToFloat": Input,
"image_tensor": Input,
"MultipleGridAnchorGenerator/Concatenate": concat_priorbox, # for 'ssd_mobilenet_v1_coco'
"Concatenate": concat_priorbox, # for other models
"concat": concat_box_loc,
"concat_1": concat_box_conf
}
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")
if model == 'ssd_mobilenet_v1_coco':
graph.find_nodes_by_name("Input")[0].input.remove("image_tensor:0")
return graph
import graphsurgeon as gs
import tensorflow as tf
Input = gs.create_node("Input",
op="Placeholder",
dtype=tf.float32,
shape=[1, 3, 512, 512])
PriorBox = gs.create_node("PriorBox",
numLayers=6,
minScale=0.2,
maxScale=0.95,
aspectRatios=[1.0, 2.0, 0.5, 3.0, 0.33],
layerVariances=[0.1,0.1,0.2,0.2],
#featureMapShapes=[19, 10, 5, 3, 2, 1])
featureMapShapes=[32, 16, 8, 4, 2, 1])
NMS = gs.create_node("NMS",
scoreThreshold=1e-8,
iouThreshold=0.6,
maxDetectionsPerClass=100,
maxTotalDetections=100,
numClasses=91,
scoreConverter="SIGMOID")
concat_priorbox = gs.create_node("concat_priorbox", dtype = tf.float32, axis = 2)
concat_box_loc = gs.create_node("concat_box_loc")
concat_box_conf = gs.create_node("concat_box_conf")
"""
namespace_plugin_map = {
"MultipleGridAnchorGenerator": PriorBox,
"Postprocessor": NMS,
"Preprocessor": Input,
"ToFloat": Input,
def ssd_mobilenet_v1_unsupported_nodes_to_plugin_nodes(ssd_graph, input_shape):
"""Makes ssd_graph TensorRT comparible using graphsurgeon.
This function takes ssd_graph, which contains graphsurgeon
DynamicGraph data structure. This structure describes frozen Tensorflow
graph, that can be modified using graphsurgeon (by deleting, adding,
replacing certain nodes). The graph is modified by removing
Tensorflow operations that are not supported by TensorRT's UffParser
and replacing them with custom layer plugin nodes.
Note: This specific implementation works only for
ssd_mobilenet_v2_coco_2018_03_29 network.
Args:
ssd_graph (gs.DynamicGraph): graph to convert
input_shape: input shape in CHW format
Returns:
gs.DynamicGraph: UffParser compatible SSD graph
"""
# Create TRT plugin nodes to replace unsupported ops in Tensorflow graph
# channels = ModelData.get_input_channels()
# height = ModelData.get_input_height()
# width = ModelData.get_input_width()
channels, height, width = input_shape
Input = gs.create_plugin_node(name="Input",
op="Placeholder",
dtype=tf.float32,
shape=[1, channels, height, width])
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],
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=91,
inputOrder=[0, 2, 1],
confSigmoid=1,
isNormalized=1
)
concat_priorbox = gs.create_node(
"concat_priorbox",
op="ConcatV2",
dtype=tf.float32,
axis=2
)
concat_box_loc = gs.create_plugin_node(
"concat_box_loc",
op="FlattenConcat_TRT",
dtype=tf.float32,
axis=1,
ignoreBatch=0
)
concat_box_conf = gs.create_plugin_node(
"concat_box_conf",
op="FlattenConcat_TRT",
dtype=tf.float32,
axis=1,
ignoreBatch=0
)
const = tf.NodeDef(name="Const", op="Const",
attr={"dtype": tf.AttrValue(type=1),
"value": tf.AttrValue(tensor=tf.make_tensor_proto([1, 1], dtype=tf.float32))})
death_list = ["strided_slice_7"]
ssd_graph.remove(ssd_graph.find_nodes_by_path(death_list))
ssd_graph.find_nodes_by_path("Shape_6")[0].input.remove("Preprocessor/sub")
# Create a mapping of namespace names -> plugin nodes.
namespace_plugin_map = {
"MultipleGridAnchorGenerator": concat_priorbox,
"Postprocessor": NMS,
"Preprocessor/map": Input,
"ToFloat": Input,
# "image_tensor": Input,
"strided_slice_6": PriorBox,
"Shape_6": const,
"Concatenate": concat_priorbox,
"concat": concat_box_loc,
"concat_1": concat_box_conf
}
for node in ssd_graph.graph_inputs:
namespace_plugin_map[node.name] = Input
# Create a new graph by collapsing namespaces
ssd_graph.collapse_namespaces(namespace_plugin_map)
# Remove the outputs, so we just have a single output node (NMS).
# If remove_exclusive_dependencies is True, the whole graph will be removed!
ssd_graph.remove(ssd_graph.graph_outputs, remove_exclusive_dependencies=False)
# Disconnect the Input node from NMS, as it expects to have only 3 inputs.
to_remove_nodes = ["Input", "Preprocessor/stack_1"]
for node in to_remove_nodes:
if node in ssd_graph.find_nodes_by_op("NMS_TRT")[0].input:
ssd_graph.find_nodes_by_op("NMS_TRT")[0].input.remove(node)
if node == "Preprocessor/stack_1":
ssd_graph.remove(node)
ssd_graph.find_nodes_by_path("Input")[0].input.remove("image_tensor:0")
return ssd_graph
def add_plugin(graph, input_dims, graph_chars=None):
graph_def = graph.as_graph_def()
if graph_chars is None:
graph_chars = converter_util.GraphCharacteristics(graph_def)
num_classes = converter_util.get_num_classes(graph_def,
graph_chars=graph_chars)
input_order = converter_util.get_NMS_input_order(graph_def,
"Postprocessor",
graph_chars=graph_chars)
if any(x == -1 for x in input_order):
print("NMS input order error: {} Aborting".format(input_order))
exit(1)
if args.debug:
print("Detected number of classes: ", num_classes)
print("Detected NMS input order: ", input_order)
assert_nodes = graph.find_nodes_by_op("Assert")
graph.remove(assert_nodes, remove_exclusive_dependencies=True)
identity_nodes = graph.find_nodes_by_op("Identity")
graph.forward_inputs(identity_nodes)
Input = gs.create_plugin_node(name="Input",
op="Placeholder",
shape=(1, ) + input_dims)
# TODO: Consider automation of parameters
PriorBox = gs.create_plugin_node(name="MultipleGridAnchorGenerator",
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],
numLayers=6)
NMS = gs.create_plugin_node(name="NMS",
op="NMS_TRT",
shareLocation=1,
varianceEncodedInTarget=0,
backgroundLabelId=0,
confidenceThreshold=0.3,
nmsThreshold=0.6,
topK=100,
keepTopK=100,
numClasses=num_classes,
inputOrder=input_order,
confSigmoid=1,
isNormalized=1)
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)
concat_priorbox = gs.create_node("concat_priorbox", op="ConcatV2", axis=2)
namespace_map = {
"MultipleGridAnchorGenerator": PriorBox,
"Preprocessor": Input,
"ToFloat": Input,
"Cast": Input,
"image_tensor": Input,
"Postprocessor": NMS,
"concat": concat_box_loc,
"concat_1": concat_box_conf,
"Concatenate": concat_priorbox,
"MultipleGridAnchorGenerator/Concatenate": concat_priorbox,
"SecondStagePostprocessor": NMS
}
graph.collapse_namespaces(namespace_map)
graph.remove(graph.graph_outputs, remove_exclusive_dependencies=False)
if graph.find_nodes_by_op("NMS_TRT"):
if "Input" in graph.find_nodes_by_op("NMS_TRT")[0].input:
graph.find_nodes_by_op("NMS_TRT")[0].input.remove("Input")
if "image_tensor:0" in graph.find_nodes_by_name("Input")[0].input:
graph.find_nodes_by_name("Input")[0].input.remove("image_tensor:0")
if "image_tensor" in graph.find_nodes_by_name("Input")[0].input:
graph.find_nodes_by_name("Input")[0].input.remove("image_tensor")
if graph.find_nodes_by_name("ToFloat_3"):
if "image_tensor:0" in graph.find_nodes_by_name("ToFloat_3")[0].input:
graph.find_nodes_by_name("ToFloat_3")[0].input.remove(
"image_tensor:0")
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",
dtype=tf.float32,
shape=[None, 3, 300, 300])
PriorBox = gs.create_plugin_node(
name="MultipleGridAnchorGenerator",
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=[32, 16, 8, 4, 2, 1],
nmsFeatureMapShapes=[19, 10, 5, 3, 2, 1],
numLayers=6)
NMS = gs.create_plugin_node(
name="NMS",
op="NMS_TRT",
inputs=['concat_box_conf', 'Squeeze', 'concat_priorbox'],
shareLocation=1,
varianceEncodedInTarget=0,
backgroundLabelId=0,
confidenceThreshold=.015,
nmsThreshold=0.4,
topK=100,
keepTopK=100,
numClasses=39, # 38 object + 1 for unknown class
inputOrder=[1, 0, 2],
confSigmoid=1,
isNormalized=1)
concat_priorbox = gs.create_node("concat_priorbox",
op="ConcatV2",
dtype=tf.float32,
inputs=['MultipleGridAnchorGenerator'],
axis=2)
concat_box_loc = gs.create_plugin_node("concat_box_loc",
op="FlattenConcat_TRT",
dtype=tf.float32,
axis=1,
ignoreBatch=0)
concat_box_conf = gs.create_plugin_node("concat_box_conf",
op="FlattenConcat_TRT",
dtype=tf.float32,
axis=1,
ignoreBatch=0)
# Create a dummy node in the 'MultipleGridAnchorGenerator' namespace.
# This is a hack for 'ssd_mobilenet_v3_large/small'...
if not any([
n.startswith('MultipleGridAnchorGenerator/')
for n in graph.node_map.keys()
]):
const = create_const_for_anchor_generator()
dummy = gs.create_node(
'MultipleGridAnchorGenerator/dummy_for_anchors',
op='Dummy', # not important here, node will be collapsed later
inputs=['const_for_anchors'])
graph.add(const)
graph.add(dummy)
namespace_plugin_map = {
"MultipleGridAnchorGenerator": PriorBox,
"Postprocessor": NMS,
"Preprocessor": Input,
"Cast": Input,
"ToFloat": Input, # Maybe replaced by Cast?
"image_tensor": Input,
'normalized_input_image_tensor': Input,
'MultipleGridAnchorGenerator/Concatenate': concat_priorbox,
"Concatenate": concat_priorbox,
"concat": concat_box_loc,
"concat_1": concat_box_conf
}
graph.collapse_namespaces(namespace_plugin_map)
if 'anchors' in [node.name for node in graph.graph_outputs]:
graph.remove(graph.graph_outputs, remove_exclusive_dependencies=False)
graph.find_nodes_by_op("NMS_TRT")[0].input.remove("Input")
if 'NMS' not in [node.name for node in graph.graph_outputs]:
graph.remove(graph.graph_outputs, remove_exclusive_dependencies=False)
if 'NMS' not in [node.name for node in graph.graph_outputs]:
# We expect 'NMS' to be one of the outputs
raise RuntimeError('bad graph_outputs')
if 'Input' in list(graph.find_nodes_by_name('NMS')[0].input):
graph.find_nodes_by_name('NMS')[0].input.remove('Input')
if 'image_tensor:0' in list(graph.find_nodes_by_name('Input')[0].input):
graph.find_nodes_by_name('Input')[0].input.remove('image_tensor:0')
graph = replace_addv2(graph)
graph = replace_fusedbnv3(graph)
return graph
import graphsurgeon as gs
import tensorflow as tf
Output = gs.create_node("dot_op_trt", op="my_matmul", dtype=tf.float32)
namespace_plugin_map = {"dot_op": Output}
def preprocess(dynamic_graph):
# Now create a new graph by collapsing namespaces
dynamic_graph.collapse_namespaces(namespace_plugin_map)
def add_plugin(cls, graph):
import tensorflow as tf
import graphsurgeon as gs
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",
dtype=tf.float32,
shape=[1, *cls.INPUT_SHAPE])
PriorBox = gs.create_plugin_node(
name="MultipleGridAnchorGenerator",
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],
numLayers=6)
NMS = gs.create_plugin_node(name="NMS",
op="NMS_TRT",
shareLocation=1,
varianceEncodedInTarget=0,
backgroundLabelId=0,
confidenceThreshold=1e-8,
nmsThreshold=cls.NMS_THRESH,
topK=100,
keepTopK=100,
numClasses=91,
inputOrder=[0, 2, 1],
confSigmoid=1,
isNormalized=1)
concat_priorbox = gs.create_node("concat_priorbox",
op="ConcatV2",
dtype=tf.float32,
axis=2)
concat_box_loc = gs.create_plugin_node("concat_box_loc",
op="FlattenConcat_TRT",
dtype=tf.float32,
axis=1,
ignoreBatch=0)
concat_box_conf = gs.create_plugin_node("concat_box_conf",
op="FlattenConcat_TRT",
dtype=tf.float32,
axis=1,
ignoreBatch=0)
namespace_plugin_map = {
"MultipleGridAnchorGenerator": PriorBox,
"Postprocessor": NMS,
"Preprocessor": Input,
"ToFloat": Input,
"image_tensor": Input,
"MultipleGridAnchorGenerator/Concatenate": concat_priorbox,
"concat": concat_box_loc,
"concat_1": concat_box_conf
}
# Create a new graph by collapsing namespaces
graph.collapse_namespaces(namespace_plugin_map)
# Remove the outputs, so we just have a single output node (NMS).
# If remove_exclusive_dependencies is True, the whole graph will be removed!
graph.remove(graph.graph_outputs, remove_exclusive_dependencies=False)
graph.find_nodes_by_op("NMS_TRT")[0].input.remove("Input")
graph.find_nodes_by_name("Input")[0].input.remove("image_tensor:0")
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
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
import graphsurgeon as gs
import tensorflow as tf
Input = gs.create_node("Input",
op="Placeholder",
dtype=tf.float32,
shape=[1, 3, 300, 300])
PriorBox = gs.create_plugin_node(name="GridAnchor",
op="GridAnchor_TRT",
numLayers=6,
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])
NMS = gs.create_plugin_node(name="NMS",
op="NMS_TRT",
shareLocation=1,
varianceEncodedInTarget=0,
backgroundLabelId=0,
confidenceThreshold=1e-8,
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=91,
inputOrder=[1, 0, 2],
confSigmoid=1,
isNormalized=1)
concat_priorbox = gs.create_node(name="concat_priorbox",
op="ConcatV2",
dtype=tf.float32,
axis=2)
concat_box_loc = gs.create_plugin_node("concat_box_loc",
op="FlattenConcat_TRT",
dtype=tf.float32,
axis=1,
ignoreBatch=0)
concat_box_conf = gs.create_plugin_node("concat_box_conf",
op="FlattenConcat_TRT",
dtype=tf.float32,
axis=1,
ignoreBatch=0)
namespace_plugin_map = {
def ssd_pipeline_to_uff(checkpoint_path,
config_path,
tmp_dir='exported_model'):
import graphsurgeon as gs
from object_detection import exporter
import tensorflow as tf
import uff
# TODO(@jwelsh): Implement by extending model builders with
# TensorRT plugin stubs. Currently, this method uses pattern
# matching which is a bit hacky and subject to fail when TF
# object detection API exporter changes. We should add object
# detection as submodule to avoid versioning incompatibilities.
config = _load_config(config_path)
frozen_graph_path = os.path.join(tmp_dir, FROZEN_GRAPH_NAME)
# get input shape
channels = 3
height = config.model.ssd.image_resizer.fixed_shape_resizer.height
width = config.model.ssd.image_resizer.fixed_shape_resizer.width
tf_config = tf.ConfigProto()
tf_config.gpu_options.allow_growth = True
# export checkpoint and config to frozen graph
with tf.Session(config=tf_config) as tf_sess:
with tf.Graph().as_default() as tf_graph:
subprocess.call(['mkdir', '-p', tmp_dir])
exporter.export_inference_graph('image_tensor',
config,
checkpoint_path,
tmp_dir,
input_shape=[1, None, None, 3])
dynamic_graph = gs.DynamicGraph(frozen_graph_path)
# remove all assert nodes
#all_assert_nodes = dynamic_graph.find_nodes_by_op("Assert")
#dynamic_graph.remove(all_assert_nodes, remove_exclusive_dependencies=True)
# forward all identity nodes
all_identity_nodes = dynamic_graph.find_nodes_by_op("Identity")
dynamic_graph.forward_inputs(all_identity_nodes)
# create input plugin
input_plugin = gs.create_plugin_node(name=TRT_INPUT_NAME,
op="Placeholder",
dtype=tf.float32,
shape=[1, height, width, channels])
# create anchor box generator
anchor_generator_config = config.model.ssd.anchor_generator.ssd_anchor_generator
box_coder_config = config.model.ssd.box_coder.faster_rcnn_box_coder
priorbox_plugin = gs.create_plugin_node(
name="priorbox",
op="GridAnchor_TRT",
minSize=anchor_generator_config.min_scale,
maxSize=anchor_generator_config.max_scale,
aspectRatios=list(anchor_generator_config.aspect_ratios),
variance=[
1.0 / box_coder_config.y_scale, 1.0 / box_coder_config.x_scale,
1.0 / box_coder_config.height_scale,
1.0 / box_coder_config.width_scale
],
featureMapShapes=_get_feature_map_shape(config),
numLayers=config.model.ssd.anchor_generator.ssd_anchor_generator.
num_layers)
# create nms plugin
nms_config = config.model.ssd.post_processing.batch_non_max_suppression
nms_plugin = gs.create_plugin_node(
name=TRT_OUTPUT_NAME,
op="NMS_TRT",
shareLocation=1,
varianceEncodedInTarget=0,
backgroundLabelId=0,
confidenceThreshold=nms_config.score_threshold,
nmsThreshold=nms_config.iou_threshold,
topK=nms_config.max_detections_per_class,
keepTopK=nms_config.max_total_detections,
numClasses=config.model.ssd.num_classes + 1, # add background
inputOrder=[1, 2, 0],
confSigmoid=1,
isNormalized=1,
scoreConverter="SIGMOID",
codeType=3)
priorbox_concat_plugin = gs.create_node("priorbox_concat",
op="ConcatV2",
dtype=tf.float32,
axis=2)
boxloc_concat_plugin = gs.create_plugin_node(
"boxloc_concat",
op="FlattenConcat_TRT_jetbot",
dtype=tf.float32,
)
boxconf_concat_plugin = gs.create_plugin_node(
"boxconf_concat",
op="FlattenConcat_TRT_jetbot",
dtype=tf.float32,
)
namespace_plugin_map = {
"MultipleGridAnchorGenerator": priorbox_plugin,
"Postprocessor": nms_plugin,
"Preprocessor": input_plugin,
"ToFloat": input_plugin,
"image_tensor": input_plugin,
"Concatenate": priorbox_concat_plugin,
"concat": boxloc_concat_plugin,
"concat_1": boxconf_concat_plugin
}
dynamic_graph.collapse_namespaces(namespace_plugin_map)
# fix name
for i, name in enumerate(
dynamic_graph.find_nodes_by_op('NMS_TRT')[0].input):
if TRT_INPUT_NAME in name:
dynamic_graph.find_nodes_by_op('NMS_TRT')[0].input.pop(i)
dynamic_graph.remove(dynamic_graph.graph_outputs,
remove_exclusive_dependencies=False)
uff_buffer = uff.from_tensorflow(dynamic_graph.as_graph_def(),
[TRT_OUTPUT_NAME])
return uff_buffer
def add_plugin(graph, model, spec):
"""add_plugin
Reference:
1. https://github.com/AastaNV/TRT_object_detection/blob/master/config/model_ssd_mobilenet_v1_coco_2018_01_28.py
2. https://github.com/AastaNV/TRT_object_detection/blob/master/config/model_ssd_mobilenet_v2_coco_2018_03_29.py
3. https://devtalk.nvidia.com/default/topic/1050465/jetson-nano/how-to-write-config-py-for-converting-ssd-mobilenetv2-to-uff-format/post/5333033/#5333033
"""
numClasses = spec['num_classes']
minSize = spec['min_size']
maxSize = spec['max_size']
inputOrder = spec['input_order']
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, ) + INPUT_DIMS)
PriorBox = gs.create_plugin_node(
name='MultipleGridAnchorGenerator',
op='GridAnchor_TRT',
minSize=minSize, # was 0.2
maxSize=maxSize, # was 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],
numLayers=6)
NMS = gs.create_plugin_node(
name='NMS',
op='NMS_TRT',
shareLocation=1,
varianceEncodedInTarget=0,
backgroundLabelId=0,
confidenceThreshold=0.3, # was 1e-8
nmsThreshold=0.6,
topK=100,
keepTopK=100,
numClasses=numClasses, # was 91
inputOrder=inputOrder,
confSigmoid=1,
isNormalized=1)
concat_priorbox = gs.create_node('concat_priorbox', op='ConcatV2', axis=2)
if trt.__version__[0] >= '7':
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)
else:
concat_box_loc = gs.create_plugin_node('concat_box_loc',
op='FlattenConcat_TRT')
concat_box_conf = gs.create_plugin_node('concat_box_conf',
op='FlattenConcat_TRT')
namespace_for_removal = [
'ToFloat',
'image_tensor',
'Preprocessor/map/TensorArrayStack_1/TensorArrayGatherV3',
]
namespace_plugin_map = {
'MultipleGridAnchorGenerator': PriorBox,
'Postprocessor': NMS,
'Preprocessor': Input,
'ToFloat': Input,
'Cast': Input, # added for models trained with tf 1.15+
'image_tensor': Input,
'MultipleGridAnchorGenerator/Concatenate':
concat_priorbox, # for 'ssd_mobilenet_v1_coco'
'Concatenate': concat_priorbox, # for other models
'concat': concat_box_loc,
'concat_1': concat_box_conf
}
graph.remove(
graph.find_nodes_by_path(
['Preprocessor/map/TensorArrayStack_1/TensorArrayGatherV3']),
remove_exclusive_dependencies=False) # for 'ssd_inception_v2_coco'
graph.collapse_namespaces(namespace_plugin_map)
graph = replace_addv2(graph)
graph = replace_fusedbnv3(graph)
if 'image_tensor:0' in graph.find_nodes_by_name('Input')[0].input:
graph.find_nodes_by_name('Input')[0].input.remove('image_tensor:0')
if 'Input' in graph.find_nodes_by_name('NMS')[0].input:
graph.find_nodes_by_name('NMS')[0].input.remove('Input')
# Remove the Squeeze to avoid "Assertion 'isPlugin(layerName)' failed"
graph.forward_inputs(
graph.find_node_inputs_by_name(graph.graph_outputs[0], 'Squeeze'))
if 'anchors' in [node.name for node in graph.graph_outputs]:
graph.remove('anchors', remove_exclusive_dependencies=False)
if len(graph.find_nodes_by_op('GridAnchor_TRT')[0].input) < 1:
graph = add_anchor_input(graph)
if 'NMS' not in [node.name for node in graph.graph_outputs]:
graph.remove(graph.graph_outputs, remove_exclusive_dependencies=False)
if 'NMS' not in [node.name for node in graph.graph_outputs]:
# We expect 'NMS' to be one of the outputs
raise RuntimeError('bad graph_outputs')
return graph
def add_plugin(graph: DynamicGraph, spec: ModelSpec):
"""add_plugin
Reference:
1. https://github.com/AastaNV/TRT_object_detection/blob/master/config/model_ssd_mobilenet_v1_coco_2018_01_28.py
2. https://github.com/AastaNV/TRT_object_detection/blob/master/config/model_ssd_mobilenet_v2_coco_2018_03_29.py
3. https://devtalk.nvidia.com/default/topic/1050465/jetson-nano/how-to-write-config-py-for-converting-ssd-mobilenetv2-to-uff-format/post/5333033/#5333033
"""
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 = create_plugin_node(name="Input",
op="Placeholder",
shape=(1, ) + spec.input_dim)
prior_box = create_plugin_node(
name="MultipleGridAnchorGenerator",
op="GridAnchor_TRT",
minSize=spec.min_size, # was 0.2
maxSize=spec.max_size, # was 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],
numLayers=6,
)
nms = create_plugin_node(
name="NMS",
op="NMS_TRT",
shareLocation=1,
varianceEncodedInTarget=0,
backgroundLabelId=0,
confidenceThreshold=0.3, # was 1e-8
nmsThreshold=0.6,
topK=100,
keepTopK=100,
numClasses=spec.num_classes,
inputOrder=spec.input_order,
confSigmoid=1,
isNormalized=1,
)
concat_priorbox = create_node("concat_priorbox", op="ConcatV2", axis=2)
concat_box_loc = create_plugin_node("concat_box_loc",
op="FlattenConcat_TRT",
axis=1,
ignoreBatch=0)
concat_box_conf = create_plugin_node("concat_box_conf",
op="FlattenConcat_TRT",
axis=1,
ignoreBatch=0)
namespace_for_removal = [
"ToFloat",
"image_tensor",
"Preprocessor/map/TensorArrayStack_1/TensorArrayGatherV3",
]
namespace_plugin_map = {
"MultipleGridAnchorGenerator": prior_box,
"Postprocessor": nms,
"Preprocessor": input,
"ToFloat": input,
"Cast": input, # added for models trained with tf 1.15+
"image_tensor": input,
"MultipleGridAnchorGenerator/Concatenate":
concat_priorbox, # for 'ssd_mobilenet_v1_coco'
"Concatenate": concat_priorbox, # for other models
"concat": concat_box_loc,
"concat_1": concat_box_conf,
}
graph.remove(
graph.find_nodes_by_path(
["Preprocessor/map/TensorArrayStack_1/TensorArrayGatherV3"]),
remove_exclusive_dependencies=False,
) # for 'ssd_inception_v2_coco'
graph.collapse_namespaces(namespace_plugin_map)
graph = replace_addv2(graph)
graph = replace_fusedbnv3(graph)
if "image_tensor:0" in graph.find_nodes_by_name("Input")[0].input:
graph.find_nodes_by_name("Input")[0].input.remove("image_tensor:0")
if "Input" in graph.find_nodes_by_name("NMS")[0].input:
graph.find_nodes_by_name("NMS")[0].input.remove("Input")
# Remove the Squeeze to avoid "Assertion 'isPlugin(layerName)' failed"
graph.forward_inputs(
graph.find_node_inputs_by_name(graph.graph_outputs[0], "Squeeze"))
if "anchors" in [node.name for node in graph.graph_outputs]:
graph.remove("anchors", remove_exclusive_dependencies=False)
if len(graph.find_nodes_by_op("GridAnchor_TRT")[0].input) < 1:
graph = add_anchor_input(graph)
if "NMS" not in [node.name for node in graph.graph_outputs]:
graph.remove(graph.graph_outputs, remove_exclusive_dependencies=False)
if "NMS" not in [node.name for node in graph.graph_outputs]:
# We expect 'NMS' to be one of the outputs
raise RuntimeError("bad graph_outputs")
return graph
def ssd_unsupported_nodes_to_plugin_nodes(ssd_graph, n_classes, input_dims,
feature_dims):
"""Makes ssd_graph TensorRT comparible using graphsurgeon.
This function takes ssd_graph, which contains graphsurgeon
DynamicGraph data structure. This structure describes frozen Tensorflow
graph, that can be modified using graphsurgeon (by deleting, adding,
replacing certain nodes). The graph is modified by removing
Tensorflow operations that are not supported by TensorRT's UffParser
and replacing them with custom layer plugin nodes.
Note: This specific implementation works only for
ssd_inception_v2_coco_2017_11_17 network.
Args:
ssd_graph (gs.DynamicGraph): graph to convert
Returns:
gs.DynamicGraph: UffParser compatible SSD graph
"""
# Remove assert nodes
all_assert_nodes = ssd_graph.find_nodes_by_op("Assert")
# Remove those nodes from the graph.
ssd_graph.remove(all_assert_nodes, remove_exclusive_dependencies=True)
# Find all identity nodes.
all_identity_nodes = ssd_graph.find_nodes_by_op("Identity")
# Forward inputs those in the graph i.e. forward their inputs.
ssd_graph.forward_inputs(all_identity_nodes)
# Create TRT plugin nodes to replace unsupported ops in Tensorflow graph
channels = input_dims[0]
height = input_dims[1]
width = input_dims[2]
nodes = ssd_graph.node_map
node_names = ssd_graph.node_map.keys()
break_now = False
class_predictor_label = 'concat_1'
box_loc_label = 'concat'
for inode in nodes['concat'].input:
if break_now:
break
for jnode in nodes[inode].input:
if break_now:
break
if 'ClassPredictor' in jnode:
class_predictor_label = 'concat'
box_loc_label = 'concat_1'
break_now = True
concat_namespace = "Concatenate"
include_anchors = False
for k in node_names:
if "MultipleGridAnchorGenerator" in k:
include_anchors = True
if "MultipleGridAnchorGenerator/Concatenate" in k:
concat_namespace = "MultipleGridAnchorGenerator/Concatenate"
# Now we need to collapse a few namespaces.
if include_anchors:
Concat = gs.create_node("concat_priorbox",
op="ConcatV2",
dtype=tf.float32,
axis=2)
Input = gs.create_plugin_node(ModelData.INPUT_NAME,
op="Placeholder",
dtype=tf.float32,
shape=[1, channels, height, width])
FlattenConcat_box_conf = gs.create_plugin_node(
"concat_box_conf",
op="FlattenConcat_TRT",
dtype=tf.float32,
)
FlattenConcat_box_loc = gs.create_plugin_node(
"concat_box_loc",
op="FlattenConcat_TRT",
dtype=tf.float32,
)
NMS = build_nms_node(numClasses=n_classes)
# Create a mapping of namespace names -> plugin nodes.
namespace_plugin_map = {
"Postprocessor": NMS,
"Preprocessor": Input,
"ToFloat": Input,
"image_tensor": Input,
box_loc_label: FlattenConcat_box_loc,
class_predictor_label: FlattenConcat_box_conf
}
# # Now create a new graph by collapsing namespaces
# ssd_graph.collapse_namespaces(namespace_plugin_map)
# # Determine the parameter for the GridAnchors
# # print(ssd_graph.as_graph_def().node)
# # print(tf.import_graph_def(ssd_graph.as_graph_def()))
# for node in ssd_graph.as_graph_def().node:
# if node.name == "BoxPredictor_2/BoxEncodingPredictor/BiasAdd":
# print(type(node))
# print(node)
# # print([dict(nodes[x].attr) for x in ssd_graph.node_map.keys() if 'BoxEncodingPredictor/BiasAdd' in x])
# # print(ssd_graph.find_nodes_by_name("BoxPredictor_2/BoxEncodingPredictor/BiasAdd"))
# exit()
if include_anchors:
GridAnchor = build_grid_anchor_node(featureMapShapes=feature_dims)
namespace_plugin_map[concat_namespace] = Concat
namespace_plugin_map["MultipleGridAnchorGenerator"] = GridAnchor
# Now create a new graph by collapsing namespaces
ssd_graph.collapse_namespaces(namespace_plugin_map)
# Remove the outputs, so we just have a single output node (NMS).
# If remove_exclusive_dependencies is True, the whole graph will be removed!
ssd_graph.remove(ssd_graph.graph_outputs,
remove_exclusive_dependencies=False)
# add in grid anchors for SSDLite
if not include_anchors:
Const = gs.create_node("Const",
op="Const",
dtype=tf.float32,
value=[128, 128])
GridAnchor = build_grid_anchor_node(inputs=[Const],
featureMapShapes=feature_dims)
Concat = gs.create_node("concat_priorbox",
inputs=[GridAnchor],
op="ConcatV2",
dtype=tf.float32,
axis=2)
ssd_graph.append(Const)
ssd_graph.append(GridAnchor)
ssd_graph.append(Concat)
NMS = build_nms_node(inputs=list(NMS.input) + ["concat_priorbox"],
numClasses=n_classes)
namespace_plugin_map = {"NMS": NMS}
ssd_graph.collapse_namespaces(namespace_plugin_map)
# For exported graphs, we need to remove the squeeze node between concat plugin and the NMS plugin.
# Downloaded graphs don't need this step. Refer to convert_ssd_v1.py
all_squeeze_nodes = ssd_graph.find_nodes_by_name("Squeeze")
# Forward inputs those in the graph i.e. forward their inputs.
ssd_graph.forward_inputs(all_squeeze_nodes)
# clean up NMS and Input nodes
actualInputOrder = []
for node in ssd_graph._internal_graphdef.node:
if node.name == "NMS":
if ModelData.INPUT_NAME in node.input:
node.input.remove(ModelData.INPUT_NAME)
for input_name in node.input:
if "loc" in input_name:
actualInputOrder.append(0)
elif "conf" in input_name:
actualInputOrder.append(1)
elif "priorbox" in input_name:
actualInputOrder.append(2)
elif node.name == ModelData.INPUT_NAME:
if "image_tensor:0" in node.input:
node.input.remove("image_tensor:0")
# NOTE: since the actual order of the NMS nodes inputs differ between versions, I'll reinsert the NMS trt op
NMS = build_nms_node(inputOrder=actualInputOrder, numClasses=n_classes)
namespace_plugin_map = {"NMS": NMS}
ssd_graph.collapse_namespaces(namespace_plugin_map)
return ssd_graph
