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)