NVIDIA / TensorRT

NVIDIA® TensorRT™ is an SDK for high-performance deep learning inference on NVIDIA GPUs. This repository contains the open source components of TensorRT.
https://developer.nvidia.com/tensorrt
Apache License 2.0
10.72k stars 2.12k forks source link

[onnx-graphsurgeon] shape inference before import_graph #3267

Closed inisis closed 1 year ago

inisis commented 1 year ago

Hi, when shape info is missing within a onnx model. calling fold_constants will fail for many cases, so I wonder if shape_infernce can be added before import_graph, a simple case can be find here test.zip

import onnx
import onnx_graphsurgeon as gs

model = onnx.load('test.onnx')

def process_graph(model, output_path):
    graph = gs.import_onnx(model)
    graph.fold_constants().cleanup().toposort()
    model_gs = gs.export_onnx(graph)
    onnx.save(model_gs, output_path)

process_graph(model, 'model_without_shape.onnx')

model_with_shape = onnx.shape_inference.infer_shapes(model, data_prop=True)
process_graph(model_with_shape, 'model_with_shape.onnx')

and a quick fix is as below, it's really hard to insert shape inference within graphsurgeon.


def import_onnx(onnx_model: "onnx.ModelProto") -> Graph:
    """
    Import an onnx-graphsurgeon Graph from the provided ONNX model.

    Args:
        onnx_model (onnx.ModelProto): The ONNX model.

    Returns:
        Graph: A corresponding onnx-graphsurgeon Graph.
    """

    onnx_model= onnx.shape_inference.infer_shapes(onnx_model, data_prop=True)

    return OnnxImporter.import_graph(
        onnx_model.graph,
        opset=OnnxImporter.get_opset(onnx_model),
        import_domains=OnnxImporter.get_import_domains(onnx_model),
        producer_name=onnx_model.producer_name,
        producer_version=onnx_model.producer_version,
    )
``

@pranavm-nvidia 
pranavm-nvidia commented 1 year ago

@inisis this is essentially what polygraphy surgeon sanitize --fold-constants does. There's also an equivalent Python API documented here

inisis commented 1 year ago

@inisis this is essentially what polygraphy surgeon sanitize --fold-constants does. There's also an equivalent Python API documented here

@pranavm-nvidia That's so great! How about add this in onnx-graphsurgeon.

ttyio commented 1 year ago

closing since this is al ready solved, thanks all!