Enhance Model Export with Custom Saver Functionality

[Han123su]

Fixes #6375 .

Description

Make the _export function focus on model conversion and saving, and add the saver parameter so that different model conversions in the future can also call the _export function, instead of making special behaviors due to specific conditions, but maintaining a consistent behavior pattern.

• Add saver_type.py with various Saver types and Handles specific model parameters using closures.
• Refactor _export function to focus on model conversion and saving, adding a saver parameter for flexibility.
• Rewrite onnx_export、ckpt_export、trt_export function to utilize the updated _export function.
• Update parameter names in save_net_with_metadata for better compatibility.

Types of changes

• Non-breaking change (fix or new feature that would not break existing functionality).
• Breaking change (fix or new feature that would cause existing functionality to change).
• New tests added to cover the changes.
• Integration tests passed locally by running ./runtests.sh -f -u --net --coverage.
• Quick tests passed locally by running ./runtests.sh --quick --unittests --disttests.
• In-line docstrings updated.
• Documentation updated, tested make html command in the docs/ folder.

[Han123su]

Hi @KumoLiu @ericspod, could you please help me review the PR or give some advice to make it better?
Thanks for taking the time to check it out.

[ericspod]

Hi @Han123su, thanks for the submission and your work but I'm sorry to say that I don't really see the value of this change. I have made a number of inline comments that I hope you'll understand are meant to be constructive and positive, it's hard to convey tone through text alone I realise!

I feel the intent here was to create a more abstract interface for saving models in different ways that we could easily extend later to include other means. I appreciate that objective but I don't see these changes doing it. The _export function can be parameterised by a callable if we wanted to separate out the actual saving functionality from the other things that are being done, but I wouldn't do this by defining new classes like CkptSaver. Since we can use callables of any sort as parameters in Python, I would use partial instead to wrap pre-defined arguments with a function and then pass that as a parameter, so there's no need to make things class-based.

In terms of common functionality there isn't much between ONNX saving and Torch/TRT saving which is why in our original code these were done separately rather than calling _export in onnx_export. By making _export more general the code that was loading weights and putting the extra_files dictionary together gets duplicated in ckpt_export and trt_export. Duplication is generally bad despite how often I and others are guilty of it, so I would avoid this and leave the code in _export.

To accomplish the goal of making this more generic and parametric, I would suggest that _export should be changed much less with the code for loading checkpoints and putting extra_files together put back in. Instead of calling save_net_with_metadata this function would instead call a callable object which would expect the same named parameters. For Torch/TRT saving that would be just save_net_with_metadata, but for ONNX this would be a partial-wrapped function or a closure function which calls onnx.save itself. This would keep most of the functionality unchanged and reduce duplication (eg. the checkpoint loading in onnx_export is duplicate and could be factored out this way).

I hope that helps and thanks again!

[Han123su]

Hi, @ericspod ,
Thank you very much for your suggestion. First, let me explain the reason for this modification:
After understanding this issue, I interpreted it as author finding that the model conversion calling _export can only save torchscript models (i.e., the models saved after the original ckpt and trt conversion). He/she mentioned that if the converter wants to use conversions such as onnx, there will be problems in the storage part. Therefore, he/she hopes to add a saver to _export to perform different storage according to the model. Thus, he/she hopes that onnx_export can also call _export like ckpt_export and trt_export, transferring all exported storage through this function.

Currently, ckpt_export and trt_export indeed have many common parameters and load weights, so doing this would be redundant. However, if it is expected that there will be a new model conversion and storage today, if _export is also called, the parameter part will not be these common parameters. Therefore, I wanted to write them internally. The saver can add these parameters when new models are stored for better management. Of course, there are no new storage models now, so it is exactly as you said, redundant and repetitive.

These are my thoughts on the modifications. Thank you for reading. Maybe I misunderstood the meaning of this issue, or perhaps the author just wanted onnx_export to call _export. Your suggestion is very reasonable, so I would like to ask if there are more direct modification suggestions for this issue, because after reading what you said, I feel that the original program call is already perfect and the most concise and direct. Thanks for spending your precious time reading this.

[ericspod]

Hi @Han123su what I would say we can do here is to do two refactorings:

1. onnx_export should call _export and use convert_to_onnx as the converter argument. The appropriate values should be passed to _export using the kwargs value. This will convert the network to ONNX format and eliminate the need to parse data and load the network as done in the body of onnx_export currently, which is a duplication of the code in _export.
2. The _export function should be changed to accept a saver: Callable argument which will replace the call to save_net_with_metadata. Where _export is used currently the save_net_with_metadata function would be passed, probably wrapped using partial to set the values for include_config_vals and append_timestamp arguments. export_onnx would instead use a different function which wraps onnx.save with the arguments that would be expected by save_net_with_metadata.

All this avoid defining new classes and introduces one new function only (the onnx.save wrapper which could be a local function inside onnx_export). Give this a thought and let me know if it makes sense, I think this is closer to what the original issue was describing would be a good solution. Thanks!

[Han123su]

Hi @Han123su what I would say we can do here is to do two refactorings:

1. onnx_export should call _export and use convert_to_onnx as the converter argument. The appropriate values should be passed to _export using the kwargs value. This will convert the network to ONNX format and eliminate the need to parse data and load the network as done in the body of onnx_export currently, which is a duplication of the code in _export.
2. The _export function should be changed to accept a saver: Callable argument which will replace the call to save_net_with_metadata. Where _export is used currently the save_net_with_metadata function would be passed, probably wrapped using partial to set the values for include_config_vals and append_timestamp arguments. export_onnx would instead use a different function which wraps onnx.save with the arguments that would be expected by save_net_with_metadata.

All this avoid defining new classes and introduces one new function only (the onnx.save wrapper which could be a local function inside onnx_export). Give this a thought and let me know if it makes sense, I think this is closer to what the original issue was describing would be a good solution. Thanks!

Ok, I understand your suggestions, and I’ll try to make the modifications accordingly and provide feedback. Thanks again for your help!