Refactor Export for Model Conversion and Saving

[Han123su]

Fixes #6375 .

Description

Changes to be made based on the previous discussion #7835.

Modify the _export function to call the saver parameter for saving different models. Rewrite the onnx_export function using the updated _export to achieve consistency in model format conversion and saving.

• Rewrite onnx_export to call _export with convert_to_onnx and appropriate kwargs.
• Add a saver: Callable parameter to _export, replacing save_net_with_metadata.
• Pass save_net_with_metadata function wrapped with partial to set parameters like include_config_vals and append_timestamp.

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.

[ericspod]

It looks much better now. I have a comment on docstrings and something optional about moving arguments around, but otherwise if all the tests pass this is good to go I feel.

[ericspod]

@binliunls This PR is response to an issue you've raised previously, would you have any comment here? Thanks!

[Han123su]

Hi, @ericspod and @binliunls, thank you very much for your suggestions. However, it seems that there might be a slight conflict in the ideas behind the two proposed changes, suggestion 1 and suggestion 2. I would like to explain my thoughts on these two modifications. Thank you!

---

• suggestion 1

ckpt_export/trt_export :

save_ts = partial(
        save_net_with_metadata,
        include_config_vals=False,
        append_timestamp=False,
)

_export :

def _export(
    ...
) -> None:
    ...
    ...
    meta_values = parser.get().pop("_meta_", None)  
    saver(net, filepath, more_extra_files=extra_files, meta_values=meta_values)    

    logger.info(f"exported to file: {filepath}.")

When calling saver, directly passing meta_values=parser.get().pop("_meta_", None) as a parameter is feasible and makes partial more concise. Although this parameter may not be required by ONNX and might be slightly inconsistent, it does not affect the saving process.

• suggestion 2

ckpt_export/trt_export :

def ckpt_export/trt_export(
    ...
    
) -> None:
    ...
    ...
    ################## Repeat part ####################
    extra_files: dict = {}
    for i in ensure_tuple(config_file_):
        filename = os.path.basename(i)
        filename, _ = os.path.splitext(filename)
        if filename in extra_files:
            raise ValueError(f"Filename part '{filename}' is given multiple times in config file list.")
        extra_files[filename] = json.dumps(ConfigParser.load_config_file(i)).encode()
    extra_files = {k + ".json": v for k, v in extra_files.items()}
    ################################################

    save_ts = partial(
        save_net_with_metadata,
        include_config_vals=False,
        append_timestamp=False,
        meta_values=parser.get().pop("_meta_", None),
        more_extra_files=extra_files,   
    )
    ...

_export :

def _export(
    ...
) -> None:
    ...
    ...
    saver(net, filepath,)   
    ...

Moving the extra_files logic from _export to ckpt_export/trt_export is a viable approach since ONNX storage indeed does not require this parameter. This method would require two parameters when calling saver, making it more consistent. However, doing so might introduce previously discussed issues and increase code redundancy (as mentioned in the "Repeat part" section above).

---

My thought is that we might perhaps adopt suggestion 1, calling the saver function and passing meta_values=parser.get().pop("_meta_", None) to make the partial simpler, without excessive repetition. Although this parameter is not required for ONNX storage, I have used the **kwargs: Any parameter in the definition of save_onnx so that it can accept other parameters. This won't actually affect the original ONNX storage.

def save_onnx(onnx_obj: Any, filename_prefix_or_stream: str , **kwargs: Any) -> None:
    onnx.save(onnx_obj, filename_prefix_or_stream)

Thank you both for reading this explanation. I hope to receive feedback again to make this even better. Thank you very much!

[KumoLiu]

Hi @Han123su, thanks for the detailed explanation. I would prefer the first proposal since parser has been passed to _export which make the meta_values more straightforward to get.
@ericspod @binliunls Could you please help review it again and share your comments. Thanks!

[ericspod]

@ericspod @binliunls Could you please help review it again and share your comments. Thanks!

I am for suggestion 1, I think it looks good now with that added and with the type issue fixed mentioned by @binliunls. Thanks!

[Han123su]

Hi @Han123su, thanks for the detailed explanation. I would prefer the first proposal since parser has been passed to _export which make the meta_values more straightforward to get.

Ok. Thank you so much for your suggestions!

[Han123su]

I am for suggestion 1, I think it looks good now with that added and with the type issue fixed mentioned by @binliunls. Thanks!

Got it, thank you for your valuable feedback!

[Han123su]

@binliunls, sorry to bother you, could you please provide some suggestions on the above discussion? Thank you!

[binliunls]

monai/bundle/scripts.py

I am OK with the suggestion 1.

[Han123su]

Hi @ericspod @KumoLiu @binliunls, I have made modifications based on the discussions above. Please help review them again, and let me know if there are any further changes needed. Thanks for taking the time to check it out.

[ericspod]

I'm good now if @binliunls @KumoLiu are. Thanks!

[KumoLiu]

Thanks for the update, looks good to me now.
Hi @binliunls, do you have any further comments here, otherwise I will try to merge this one.

[binliunls]

Thanks for the update, looks good to me now. Hi @binliunls, do you have any further comments here, otherwise I will try to merge this one.

LGTM.

[KumoLiu]

/build