From a769c84bdef9c43b935eb2929b1d179d28fc552c Mon Sep 17 00:00:00 2001 From: Han123su Date: Sat, 20 Jul 2024 20:27:13 +0800 Subject: [PATCH 01/10] Modify _export with Saver and onnx_export Signed-off-by: Han123su --- monai/bundle/scripts.py | 57 +++++++++++++++++++++++++++++------------ 1 file changed, 40 insertions(+), 17 deletions(-) diff --git a/monai/bundle/scripts.py b/monai/bundle/scripts.py index 56146546e8..42cdb937e8 100644 --- a/monai/bundle/scripts.py +++ b/monai/bundle/scripts.py @@ -22,7 +22,7 @@ from pydoc import locate from shutil import copyfile from textwrap import dedent -from typing import Any, Callable +from typing import Any, Callable, IO import torch from torch.cuda import is_available @@ -35,6 +35,7 @@ from monai.bundle.workflows import BundleWorkflow, ConfigWorkflow from monai.config import IgniteInfo, PathLike from monai.data import load_net_with_metadata, save_net_with_metadata +from functools import partial from monai.networks import ( convert_to_onnx, convert_to_torchscript, @@ -1159,6 +1160,7 @@ def verify_net_in_out( def _export( converter: Callable, + saver: Callable, parser: ConfigParser, net_id: str, filepath: str, @@ -1173,6 +1175,8 @@ def _export( Args: converter: a callable object that takes a torch.nn.module and kwargs as input and converts the module to another type. + saver: a callable object that takes the converted model and a filepath as input and + saves the model to the specified location. parser: a ConfigParser of the bundle to be converted. net_id: ID name of the network component in the parser, it must be `torch.nn.Module`. filepath: filepath to export, if filename has no extension, it becomes `.ts`. @@ -1212,14 +1216,12 @@ def _export( # add .json extension to all extra files which are always encoded as JSON extra_files = {k + ".json": v for k, v in extra_files.items()} - save_net_with_metadata( - jit_obj=net, - filename_prefix_or_stream=filepath, - include_config_vals=False, - append_timestamp=False, - meta_values=parser.get().pop("_meta_", None), - more_extra_files=extra_files, + saver( + net, + filepath, + more_extra_files = extra_files, ) + logger.info(f"exported to file: {filepath}.") @@ -1318,17 +1320,27 @@ def onnx_export( input_shape_ = _get_fake_input_shape(parser=parser) inputs_ = [torch.rand(input_shape_)] - net = parser.get_parsed_content(net_id_) - if has_ignite: - # here we use ignite Checkpoint to support nested weights and be compatible with MONAI CheckpointSaver - Checkpoint.load_objects(to_load={key_in_ckpt_: net}, checkpoint=ckpt_file_) - else: - ckpt = torch.load(ckpt_file_) - copy_model_state(dst=net, src=ckpt if key_in_ckpt_ == "" else ckpt[key_in_ckpt_]) converter_kwargs_.update({"inputs": inputs_, "use_trace": use_trace_}) - onnx_model = convert_to_onnx(model=net, **converter_kwargs_) - onnx.save(onnx_model, filepath_) + + def save_onnx( + onnx_obj: torch.nn.Module, + filename_prefix_or_stream: str | IO[Any], + **kwargs: Any, + ) -> None: + onnx.save(onnx_obj, filename_prefix_or_stream) + + _export( + convert_to_onnx, + save_onnx, + parser, + net_id=net_id_, + filepath=filepath_, + ckpt_file=ckpt_file_, + config_file=config_file_, + key_in_ckpt=key_in_ckpt_, + **converter_kwargs_, + ) def ckpt_export( @@ -1449,8 +1461,14 @@ def ckpt_export( converter_kwargs_.update({"inputs": inputs_, "use_trace": use_trace_}) # Use the given converter to convert a model and save with metadata, config content + + save_ts = partial(save_net_with_metadata, include_config_vals=False, + append_timestamp=False, + meta_values=parser.get().pop("_meta_", None)) + _export( convert_to_torchscript, + save_ts, parser, net_id=net_id_, filepath=filepath_, @@ -1620,8 +1638,13 @@ def trt_export( } converter_kwargs_.update(trt_api_parameters) + save_ts = partial(save_net_with_metadata, include_config_vals=False, + append_timestamp=False, + meta_values=parser.get().pop("_meta_", None)) + _export( convert_to_trt, + save_ts, parser, net_id=net_id_, filepath=filepath_, From c70f24e0008e7ba00114e8d0fdab12f33473cc4f Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 20 Jul 2024 13:23:01 +0000 Subject: [PATCH 02/10] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- monai/bundle/scripts.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/monai/bundle/scripts.py b/monai/bundle/scripts.py index 42cdb937e8..3db058fee5 100644 --- a/monai/bundle/scripts.py +++ b/monai/bundle/scripts.py @@ -1221,7 +1221,7 @@ def _export( filepath, more_extra_files = extra_files, ) - + logger.info(f"exported to file: {filepath}.") @@ -1329,7 +1329,7 @@ def save_onnx( **kwargs: Any, ) -> None: onnx.save(onnx_obj, filename_prefix_or_stream) - + _export( convert_to_onnx, save_onnx, @@ -1465,7 +1465,7 @@ def ckpt_export( save_ts = partial(save_net_with_metadata, include_config_vals=False, append_timestamp=False, meta_values=parser.get().pop("_meta_", None)) - + _export( convert_to_torchscript, save_ts, From 0c984d91b8d4f0efce302156bf48956a81b52d1f Mon Sep 17 00:00:00 2001 From: Han123su Date: Sat, 20 Jul 2024 21:32:17 +0800 Subject: [PATCH 03/10] ./runtests.sh --autofix Signed-off-by: Han123su --- monai/bundle/scripts.py | 36 +++++++++++++++++------------------- 1 file changed, 17 insertions(+), 19 deletions(-) diff --git a/monai/bundle/scripts.py b/monai/bundle/scripts.py index 42cdb937e8..8397c09f00 100644 --- a/monai/bundle/scripts.py +++ b/monai/bundle/scripts.py @@ -18,11 +18,12 @@ import warnings import zipfile from collections.abc import Mapping, Sequence +from functools import partial from pathlib import Path from pydoc import locate from shutil import copyfile from textwrap import dedent -from typing import Any, Callable, IO +from typing import IO, Any, Callable import torch from torch.cuda import is_available @@ -35,7 +36,6 @@ from monai.bundle.workflows import BundleWorkflow, ConfigWorkflow from monai.config import IgniteInfo, PathLike from monai.data import load_net_with_metadata, save_net_with_metadata -from functools import partial from monai.networks import ( convert_to_onnx, convert_to_torchscript, @@ -1216,12 +1216,8 @@ def _export( # add .json extension to all extra files which are always encoded as JSON extra_files = {k + ".json": v for k, v in extra_files.items()} - saver( - net, - filepath, - more_extra_files = extra_files, - ) - + saver(net, filepath, more_extra_files=extra_files) + logger.info(f"exported to file: {filepath}.") @@ -1323,13 +1319,9 @@ def onnx_export( converter_kwargs_.update({"inputs": inputs_, "use_trace": use_trace_}) - def save_onnx( - onnx_obj: torch.nn.Module, - filename_prefix_or_stream: str | IO[Any], - **kwargs: Any, - ) -> None: + def save_onnx(onnx_obj: torch.nn.Module, filename_prefix_or_stream: str | IO[Any], **kwargs: Any) -> None: onnx.save(onnx_obj, filename_prefix_or_stream) - + _export( convert_to_onnx, save_onnx, @@ -1462,10 +1454,13 @@ def ckpt_export( converter_kwargs_.update({"inputs": inputs_, "use_trace": use_trace_}) # Use the given converter to convert a model and save with metadata, config content - save_ts = partial(save_net_with_metadata, include_config_vals=False, + save_ts = partial( + save_net_with_metadata, + include_config_vals=False, append_timestamp=False, - meta_values=parser.get().pop("_meta_", None)) - + meta_values=parser.get().pop("_meta_", None), + ) + _export( convert_to_torchscript, save_ts, @@ -1638,9 +1633,12 @@ def trt_export( } converter_kwargs_.update(trt_api_parameters) - save_ts = partial(save_net_with_metadata, include_config_vals=False, + save_ts = partial( + save_net_with_metadata, + include_config_vals=False, append_timestamp=False, - meta_values=parser.get().pop("_meta_", None)) + meta_values=parser.get().pop("_meta_", None), + ) _export( convert_to_trt, From b62c4d5eedb3db52ab252421fb5876ab97f18beb Mon Sep 17 00:00:00 2001 From: Han123su Date: Sat, 20 Jul 2024 21:52:26 +0800 Subject: [PATCH 04/10] modify Signed-off-by: Han123su --- monai/bundle/scripts.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/monai/bundle/scripts.py b/monai/bundle/scripts.py index 8397c09f00..d4410b3c0a 100644 --- a/monai/bundle/scripts.py +++ b/monai/bundle/scripts.py @@ -1216,7 +1216,11 @@ def _export( # add .json extension to all extra files which are always encoded as JSON extra_files = {k + ".json": v for k, v in extra_files.items()} - saver(net, filepath, more_extra_files=extra_files) + saver( + net, + filepath, + more_extra_files = extra_files, + ) logger.info(f"exported to file: {filepath}.") @@ -1458,8 +1462,12 @@ def ckpt_export( save_net_with_metadata, include_config_vals=False, append_timestamp=False, +<<<<<<< HEAD meta_values=parser.get().pop("_meta_", None), ) +======= + meta_values=parser.get().pop("_meta_", None)) +>>>>>>> c70f24e0... [pre-commit.ci] auto fixes from pre-commit.com hooks _export( convert_to_torchscript, From bf5b04d5c7f0e8b16879db9b2503c0857503edf4 Mon Sep 17 00:00:00 2001 From: Han123su Date: Sat, 20 Jul 2024 22:02:19 +0800 Subject: [PATCH 05/10] ./runtests.sh --autofix Signed-off-by: Han123su --- monai/bundle/scripts.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/monai/bundle/scripts.py b/monai/bundle/scripts.py index 515c9efe49..8397c09f00 100644 --- a/monai/bundle/scripts.py +++ b/monai/bundle/scripts.py @@ -1216,11 +1216,7 @@ def _export( # add .json extension to all extra files which are always encoded as JSON extra_files = {k + ".json": v for k, v in extra_files.items()} - saver( - net, - filepath, - more_extra_files = extra_files, - ) + saver(net, filepath, more_extra_files=extra_files) logger.info(f"exported to file: {filepath}.") From 5798d72f92929c09561e607a5972e5616b876666 Mon Sep 17 00:00:00 2001 From: Han123su Date: Sat, 20 Jul 2024 23:38:40 +0800 Subject: [PATCH 06/10] modify save_onnx Signed-off-by: Han123su --- monai/bundle/scripts.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/monai/bundle/scripts.py b/monai/bundle/scripts.py index 8397c09f00..e78037f82f 100644 --- a/monai/bundle/scripts.py +++ b/monai/bundle/scripts.py @@ -1319,7 +1319,7 @@ def onnx_export( converter_kwargs_.update({"inputs": inputs_, "use_trace": use_trace_}) - def save_onnx(onnx_obj: torch.nn.Module, filename_prefix_or_stream: str | IO[Any], **kwargs: Any) -> None: + def save_onnx(onnx_obj: Any, filename_prefix_or_stream: str | IO[Any], **kwargs: Any) -> None: onnx.save(onnx_obj, filename_prefix_or_stream) _export( From 84355a2b83ee6ce6f31fccd09353509b8bbb4864 Mon Sep 17 00:00:00 2001 From: Han123su Date: Thu, 8 Aug 2024 00:27:07 +0800 Subject: [PATCH 07/10] saver include argument meta_values Signed-off-by: Han123su --- monai/bundle/scripts.py | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/monai/bundle/scripts.py b/monai/bundle/scripts.py index e78037f82f..c75c236777 100644 --- a/monai/bundle/scripts.py +++ b/monai/bundle/scripts.py @@ -1175,9 +1175,8 @@ def _export( Args: converter: a callable object that takes a torch.nn.module and kwargs as input and converts the module to another type. - saver: a callable object that takes the converted model and a filepath as input and - saves the model to the specified location. - parser: a ConfigParser of the bundle to be converted. + saver: a callable object that accepts the converted model to save, a filepath to save to, meta values + (extracted from the parser), and a dictionary of extra JSON files (name -> contents) as input. net_id: ID name of the network component in the parser, it must be `torch.nn.Module`. filepath: filepath to export, if filename has no extension, it becomes `.ts`. ckpt_file: filepath of the model checkpoint to load. @@ -1216,7 +1215,8 @@ def _export( # add .json extension to all extra files which are always encoded as JSON extra_files = {k + ".json": v for k, v in extra_files.items()} - saver(net, filepath, more_extra_files=extra_files) + meta_values = parser.get().pop("_meta_", None) + saver(net, filepath, meta_values=meta_values, more_extra_files=extra_files) logger.info(f"exported to file: {filepath}.") @@ -1319,7 +1319,7 @@ def onnx_export( converter_kwargs_.update({"inputs": inputs_, "use_trace": use_trace_}) - def save_onnx(onnx_obj: Any, filename_prefix_or_stream: str | IO[Any], **kwargs: Any) -> None: + def save_onnx(onnx_obj: Any, filename_prefix_or_stream: str, **kwargs: Any) -> None: onnx.save(onnx_obj, filename_prefix_or_stream) _export( @@ -1458,7 +1458,6 @@ def ckpt_export( save_net_with_metadata, include_config_vals=False, append_timestamp=False, - meta_values=parser.get().pop("_meta_", None), ) _export( @@ -1637,7 +1636,6 @@ def trt_export( save_net_with_metadata, include_config_vals=False, append_timestamp=False, - meta_values=parser.get().pop("_meta_", None), ) _export( From e84dbb3b35cc87ff51f6564f1fd55bd730742539 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 11 Aug 2024 09:41:43 +0000 Subject: [PATCH 08/10] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- monai/bundle/scripts.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/monai/bundle/scripts.py b/monai/bundle/scripts.py index 41a958fc62..5200d9375d 100644 --- a/monai/bundle/scripts.py +++ b/monai/bundle/scripts.py @@ -23,7 +23,7 @@ from pydoc import locate from shutil import copyfile from textwrap import dedent -from typing import IO, Any, Callable +from typing import Any, Callable import torch from torch.cuda import is_available From 0959b08e680ef3713341edbeff8e19eee3c69c57 Mon Sep 17 00:00:00 2001 From: Han123su Date: Sun, 11 Aug 2024 17:49:44 +0800 Subject: [PATCH 09/10] ./runtests.sh --autofix Signed-off-by: Han123su --- monai/bundle/scripts.py | 12 ++---------- 1 file changed, 2 insertions(+), 10 deletions(-) diff --git a/monai/bundle/scripts.py b/monai/bundle/scripts.py index 41a958fc62..167fb35140 100644 --- a/monai/bundle/scripts.py +++ b/monai/bundle/scripts.py @@ -1549,11 +1549,7 @@ def ckpt_export( converter_kwargs_.update({"inputs": inputs_, "use_trace": use_trace_}) # Use the given converter to convert a model and save with metadata, config content - save_ts = partial( - save_net_with_metadata, - include_config_vals=False, - append_timestamp=False, - ) + save_ts = partial(save_net_with_metadata, include_config_vals=False, append_timestamp=False) _export( convert_to_torchscript, @@ -1727,11 +1723,7 @@ def trt_export( } converter_kwargs_.update(trt_api_parameters) - save_ts = partial( - save_net_with_metadata, - include_config_vals=False, - append_timestamp=False, - ) + save_ts = partial(save_net_with_metadata, include_config_vals=False, append_timestamp=False) _export( convert_to_trt, From 379728b7d375fcc9881b5dd97263515fd38eea22 Mon Sep 17 00:00:00 2001 From: Han123su Date: Sun, 11 Aug 2024 18:19:11 +0800 Subject: [PATCH 10/10] restore parser args Signed-off-by: Han123su --- monai/bundle/scripts.py | 1 + 1 file changed, 1 insertion(+) diff --git a/monai/bundle/scripts.py b/monai/bundle/scripts.py index b2d00dfc1a..142a366669 100644 --- a/monai/bundle/scripts.py +++ b/monai/bundle/scripts.py @@ -1272,6 +1272,7 @@ def _export( converts the module to another type. saver: a callable object that accepts the converted model to save, a filepath to save to, meta values (extracted from the parser), and a dictionary of extra JSON files (name -> contents) as input. + parser: a ConfigParser of the bundle to be converted. net_id: ID name of the network component in the parser, it must be `torch.nn.Module`. filepath: filepath to export, if filename has no extension, it becomes `.ts`. ckpt_file: filepath of the model checkpoint to load.