As shown in the polygraphy run tutorial TensorRT/tools/Polygraphy/examples/cli/run/08_adding_precision_constraints at main · NVIDIA/TensorRT · GitHub we can use a TensorRT network postprocessing script to apply precisions on the parsed network by adding an option with --trt-network-postprocess-script ./add_constraints.py .
However, it seems not flexible enough. I would like to call the postprocessing script with an option parameter config when calling the postprocess func to add constrain precisions to the network. How should I add the expansion to the polygraphy run tool?
"""
Postprocessing script to add precision constraints to a TensorRT network.
"""
import tensorrt as trt
def postprocess(network, config):
"""
Traverses the parsed network and constrains precisions
for specific layers to FP32.
Args:
network (trt.INetworkDefinition): The network to modify.
config : a List of target layers which are expected to add constrains precisions
Returns:
None
"""
for layer in network:
# Set computation precision for layers declared in the `config` to FP32
if layer.name in config:
layer.precision = trt.float32
# Set the output precision for the layer declared in the `config` to FP32. Without this,
# the intermediate output data of the Add may be stored as FP16 even
# though the computation itself is performed in FP32.
if layer.name in config:
layer.set_output_type(0, trt.float32)