Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Improve ckpt_export #6965

Merged
merged 14 commits into from
Sep 11, 2023
29 changes: 24 additions & 5 deletions monai/bundle/scripts.py
Original file line number Diff line number Diff line change
Expand Up @@ -1214,9 +1214,13 @@ def ckpt_export(

Args:
net_id: ID name of the network component in the config, it must be `torch.nn.Module`.
Default to "network_def".
filepath: filepath to export, if filename has no extension it becomes `.ts`.
Default to "models/model.ts" under "os.getcwd()" if `bundle_root` is not specified.
ckpt_file: filepath of the model checkpoint to load.
Default to "models/model.pt" under "os.getcwd()" if `bundle_root` is not specified.
meta_file: filepath of the metadata file, if it is a list of file paths, the content of them will be merged.
Default to "configs/metadata.json" under "os.getcwd()" if `bundle_root` is not specified.
config_file: filepath of the config file to save in TorchScript model and extract network information,
the saved key in the TorchScript model is the config filename without extension, and the saved config
value is always serialized in JSON format no matter the original file format is JSON or YAML.
Expand Down Expand Up @@ -1250,9 +1254,10 @@ def ckpt_export(
)
_log_input_summary(tag="ckpt_export", args=_args)
(
config_file_,
filepath_,
ckpt_file_,
config_file_,
bundle_root_,
net_id_,
meta_file_,
key_in_ckpt_,
Expand All @@ -1261,10 +1266,11 @@ def ckpt_export(
converter_kwargs_,
) = _pop_args(
_args,
"filepath",
"ckpt_file",
"config_file",
net_id="",
filepath=None,
ckpt_file=None,
bundle_root=os.getcwd(),
net_id=None,
meta_file=None,
key_in_ckpt="",
use_trace=False,
Expand All @@ -1275,9 +1281,22 @@ def ckpt_export(
parser = ConfigParser()

parser.read_config(f=config_file_)
if meta_file_ is not None:
meta_file_ = (
os.path.join(bundle_root_, "configs", "metadata.json") if meta_file_ is None else meta_file_
)
filepath_ = os.path.join(bundle_root_, "models", "model.ts") if filepath_ is None else filepath_
ckpt_file_ = os.path.join(bundle_root_, "models", "model.pt") if ckpt_file_ is None else ckpt_file_
if not os.path.exists(ckpt_file_):
raise FileNotFoundError(f"ckpt_file in {ckpt_file_} does not exist, please specify it.")
KumoLiu marked this conversation as resolved.
Show resolved Hide resolved
if os.path.exists(meta_file_):
parser.read_meta(f=meta_file_)

net_id_ = "network_def" if net_id_ is None else net_id_
try:
parser.get_parsed_content(net_id_)
except ValueError as e:
raise ValueError(f"net_id: network_def in {config_file_} does not exist.") from e
KumoLiu marked this conversation as resolved.
Show resolved Hide resolved

# the rest key-values in the _args are to override config content
for k, v in _args.items():
parser[k] = v
Expand Down
23 changes: 23 additions & 0 deletions tests/test_bundle_ckpt_export.py
Original file line number Diff line number Diff line change
Expand Up @@ -75,6 +75,29 @@ def test_export(self, key_in_ckpt, use_trace):
self.assertTrue("meta_file" in json.loads(extra_files["def_args.json"]))
self.assertTrue("network_def" in json.loads(extra_files["inference.json"]))

@parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3])
def test_default_value(self, key_in_ckpt, use_trace):
config_file = os.path.join(os.path.dirname(__file__), "testing_data", "inference.json")
with tempfile.TemporaryDirectory() as tempdir:
def_args = {"meta_file": "will be replaced by `meta_file` arg"}
def_args_file = os.path.join(tempdir, "def_args.yaml")
ckpt_file = os.path.join(tempdir, "models/model.pt")
ts_file = os.path.join(tempdir, "models/model.ts")

parser = ConfigParser()
parser.export_config_file(config=def_args, filepath=def_args_file)
parser.read_config(config_file)
net = parser.get_parsed_content("network_def")
save_state(src=net if key_in_ckpt == "" else {key_in_ckpt: net}, path=ckpt_file)

# check with default value
cmd = ["coverage", "run", "-m", "monai.bundle", "ckpt_export", "--key_in_ckpt", key_in_ckpt]
cmd += ["--config_file", config_file, "--bundle_root", tempdir]
if use_trace == "True":
cmd += ["--use_trace", use_trace, "--input_shape", "[1, 1, 96, 96, 96]"]
command_line_tests(cmd)
self.assertTrue(os.path.exists(ts_file))


if __name__ == "__main__":
unittest.main()
Loading