Using a Customed Network Postprocessing Script postprocess(network, config) to Constrain Precisions

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)

Hi,

Sorry for the delayed response, we are looking into this issue and get back to you soon.

Thank you.

Hi,

Hope your query is answered here Using a Customed Network Postprocessing Script postprocess(network, config) to Constrain Precisions · Issue #2508 · NVIDIA/TensorRT · GitHub.

Thank you.