Support for img norm in graph during onnx export

[nick-clarifai]

trafficstars

Motivation

The purpose of this PR is to provide an option that does the image norm step of inference within the graph itself during the export to an ONNX model.

Modification

mmseg/datasets/pipelines/transforms.py:

• an init flag is added which causes the image norm to be skipped during __call__ and causes the flag to be passed along in the img_norm_cfg

mmseg/models/segmentors/base.py

• forward modified to do image normalization to the input if img_norm_cfg['normalize_in_graph'] is True

tools/pytorch2onnx.py:

• command line arg added
• normalize step in test pipeline is modified to skip normalization but still add img_norm_cfg
• input is cloned in verification step so the pre-image norm input can be compared

Use cases

This would allow one to have a totally contained model that doesn't require extra python code for inference so that the ONNX model could easily be used in a setting such as NVIDIA Triton.

[CLAassistant]

All committers have signed the CLA.

[codecov[bot]]

Codecov Report

Attention: Patch coverage is 27.27273% with 8 lines in your changes missing coverage. Please review.

Project coverage is 89.98%. Comparing base (9975c67) to head (f62cf69). Report is 257 commits behind head on master.

Files with missing lines	Patch %	Lines
mmseg/models/segmentors/base.py	0.00%	6 Missing and 1 partial :warning:
mmseg/datasets/pipelines/transforms.py	75.00%	0 Missing and 1 partial :warning:

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #1027      +/-   ##
==========================================
- Coverage   90.11%   89.98%   -0.14%     
==========================================
  Files         125      125              
  Lines        7262     7277      +15     
  Branches     1206     1210       +4     
==========================================
+ Hits         6544     6548       +4     
- Misses        515      524       +9     
- Partials      203      205       +2     

Flag	Coverage Δ
unittests	89.98% <27.27%> (-0.14%)	:arrow_down:

Flags with carried forward coverage won't be shown. Click here to find out more.

:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.

[Junjun2016]

Hi @nick-clarifai Thanks for your contribution. We will review it in the near future.

[RockeyCoss]

Please merge the master branch into your branch, thank you.

[Junjun2016]

Hi @RunningLeon Please review it.

[RunningLeon]

@nick-clarifai Hi, thanks for your PR. If you add --normalize-in-graph in pytorch2onnx.py, then you might need to support it in downstream tool onnx2tensorrt.py as well.

[nick-clarifai]

@nick-clarifai Hi, thanks for your PR. If you add --normalize-in-graph in pytorch2onnx.py, then you might need to support it in downstream tool onnx2tensorrt.py as well.

Yeah this makes sense. The changes I can see being necessary would be allowing image normalization in the preprocessing step to be skipped and adding a command line arg to skip normalization? I'm happy to make these changes if we're on the same page. Let me know if there's something else I'm forgetting there.

[RunningLeon]

@nick-clarifai Hi, thanks for your PR. If you add --normalize-in-graph in pytorch2onnx.py, then you might need to support it in downstream tool onnx2tensorrt.py as well.

Yeah this makes sense. The changes I can see being necessary would be allowing image normalization in the preprocessing step to be skipped and adding a command line arg to skip normalization? I'm happy to make these changes if we're on the same page. Let me know if there's something else I'm forgetting there.

Yes. --normalize-in-graph should be added to onnx2tensorrt.py as well.

[nick-clarifai]

Thanks for the feedback. onnx2tensorrt.py now has an option to skip image normalization in the preprocessing step. I didn't think it made sense to name the arg --normalize-in-graph because the tensorrt model will only do so if the input onnx model does normalize in graph, so I named the arg --skip-normalize instead. Let me know if you suggest any more changes.

[RunningLeon]

Thanks for the feedback. onnx2tensorrt.py now has an option to skip image normalization in the preprocessing step. I didn't think it made sense to name the arg --normalize-in-graph because the tensorrt model will only do so if the input onnx model does normalize in graph, so I named the arg --skip-normalize instead. Let me know if you suggest any more changes.

Actually, skip-normalize works for me.

[nick-clarifai]

Is there something else that should be done before this could be merged? Let me know if there's anything I can do.

[RunningLeon]

Is there something else that should be done before this could be merged? Let me know if there's anything I can do.

@nick-clarifai Hi, sorry for the late reply. After resolve the rest comments, this PR is ready to go merge for me.

[nick-clarifai]

Is there something else that should be done before this could be merged? Let me know if there's anything I can do.

@nick-clarifai Hi, sorry for the late reply. After resolve the rest comments, this PR is ready to go merge for me.

Thanks for the feedback @RunningLeon. I've made the requested changes.

[MengzhangLI]

Hi, @nick-clarifai nick, sorry for late reply.

This pr is closely to merge, could you please fix unit test problem? Your modification is not coveraged by unit test.

Best,

[nick-clarifai]

Hi, @nick-clarifai nick, sorry for late reply.

This pr is closely to merge, could you please fix unit test problem? Your modification is not coveraged by unit test.

Best,

Sure, can you help me decide which is the best test to write? Should I write a test that runs the code from pytorch2onnx.py with the verification enabled? Or should I do something like add a test to tests/test_models/test_segmentors/test_encoder_decoder.py that makes sure the added code in mmseg/models/segmentors/base.py is run?

[OpenMMLab-Assistant005]

Hi @nick-clarifai ！We are grateful for your efforts in helping improve this open-source project during your personal time. Welcome to join OpenMMLab Special Interest Group (SIG) private channel on Discord, where you can share your experiences, ideas, and build connections with like-minded peers. To join the SIG channel, simply message moderator— OpenMMLab on Discord or briefly share your open-source contributions in the #introductions channel and we will assist you. Look forward to seeing you there! Join us ：https://discord.gg/UjgXkPWNqA If you have a WeChat account，welcome to join our community on WeChat. You can add our assistant ：openmmlabwx. Please add "mmsig + Github ID" as a remark when adding friends：） Thank you again for your contribution❤