Repository navigation
Add in-tree isaac_deploy_trt and optimize_graph TensorRT rewrite. - #16
northernLights6 wants to merge 5 commits into
Conversation
f58bfc3 to
6fc1a96
Compare
Vendors TensorRT ONNX surgery under packages/isaac_deploy_trt, exposes optimize_graph(trt_compatible=True) after compile_graph, and stamps YAML checksums plus parameters.tensorrt_compatible on exported ONNX models.
6fc1a96 to
c15ff43
Compare
Signed-off-by: Tao Zhang <tazhang@nvidia.com>
frlai
left a comment
frlai
left a comment
There was a problem hiding this comment.
there are a few major issues that will cause bugs if not taken care of. additionally, I think a major design choice is to only enable optimizations inside the optimize_graph() leapp lifecycle function and not allow optimizations the export cycle. this should help us better separate concerns and keep all the trt related stuff inside the package.
Finally, it would be good if we can revisit the optimize_graph() api. we might want to be able to add more options to it for future expandability.
There was a problem hiding this comment.
lets remove this file. Not sure if it is needed
There was a problem hiding this comment.
for the generation tool. I recall it requires internal docs/repos right? That means a user won't be able to run this either way. Just for my information, would it be a good idea to not release the scripts folder at all?
|
|
||
| ``optimization`` | ||
| Present on ONNX backends unless | ||
| ``backend_params={"annotate_tensorrt_compatible": False}``. |
There was a problem hiding this comment.
lets not have this option. we can centralize all the optimizations in the optimize_graph api
| interop). Use this only to test whether those patches are involved in | ||
| a failure. | ||
|
|
||
| TensorRT catalog and optimize |
There was a problem hiding this comment.
this deserves its own docs page. can you move make a new page called something like optimizations? keep it next to debugging.
| ``opset_version`` | ||
| Present on ONNX backends. The ONNX opset used at export. | ||
|
|
||
| ``optimization`` |
There was a problem hiding this comment.
I don't think this should go into parameters, optimization should be centralized in optimize graph. when building the onnx file we should not be able to replace with tensorrt compatible ops.
| _get_logger().section("Graph optimization: trt_compatible") | ||
| result = rewrite_leapp_export_for_tensorrt(save_path) | ||
| rewritten = [model.filename for model in result.models if model.rewritten] | ||
| if rewritten: |
There was a problem hiding this comment.
if rewritten we should also update the yaml and the model hash.
previously we kept that in the onnx backend but doing it here is probably cleaner.
| return True | ||
|
|
||
|
|
||
| def optimize_graph(trt_compatible=False): |
There was a problem hiding this comment.
make sure this handles the case where one or some models aren't onnx.
there should also be some sort of way to only select a few models to do the optimizations for?
maybe something like
optimize_graph(validate_trt=['model 1', 'model 2'])
or also
optimize_graph(validate_trt=['all'])
There was a problem hiding this comment.
considering future expansions, I am also considering maybe we have a more complex system for more fine grained control?
perhaps
leapp.apply_optimizations('model 1' : [leapp.trt_compatibility(), leapp.quantization('fp16')], 'model_2' : [...])
just a thought.
There was a problem hiding this comment.
agent commented that if we install this with pip, somehow we are missing the catalogs. Might need to explicitly package them inside the pyproject.
| "Pass trt_compatible=True to optimize ONNX for TensorRT." | ||
| ) | ||
|
|
||
| return applied |
There was a problem hiding this comment.
since we change numerical behavior, it would be best to run a validation run with configurable tolerance just like the compile graph if any changes were applied to make sure the export with the changes is within tolerance.
This should be similar to how compile graph does it.
|
|
||
| graph = load_graph(src) | ||
| rewritten_nodes = apply_tensorrt_rewrites(graph, surgeon_ops=surgeon_ops) | ||
| save_graph(graph, dst) |
There was a problem hiding this comment.
Agent noticed that:
When the input ONNX uses external weights and this rewrite branch runs, saving into a different directory leaves the output referencing external data that hasn’t been copied. For example, input/model.onnx references model.onnx.data, but the output directory receives only model.onnx; loading it then fails because the weights file is missing.
I think the solution is we need to check if the onnx file has external data and also copy that over.
Vendors TensorRT ONNX surgery under packages/isaac_deploy_trt, exposes optimize_graph(trt_compatible=True) after compile_graph, and stamps YAML checksums plus parameters.tensorrt_compatible on exported ONNX models.