HongChuang/ultralytics-ascend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 3

Add export args to model metadata (#18472)

Browse source 
Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Glenn Jocher 2025-01-02 21:26:05 +01:00 • committed by GitHub
parent 990424cb63
commit 1f1f355a9e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 3 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

4
ultralytics/engine/exporter.py
View file

@ -250,7 +250,8 @@ class Exporter:
self.device = select_device("cpu" if self.args.device is None else self.args.device) self.device = select_device("cpu" if self.args.device is None else self.args.device)
# Argument compatibility checks # Argument compatibility checks
validate_args(fmt, self.args, fmts_dict["Arguments"][flags.index(True) + 1]) fmt_keys = fmts_dict["Arguments"][flags.index(True) + 1]
validate_args(fmt, self.args, fmt_keys)
if imx and not self.args.int8: if imx and not self.args.int8:
LOGGER.warning("WARNING ⚠️ IMX only supports int8 export, setting int8=True.") LOGGER.warning("WARNING ⚠️ IMX only supports int8 export, setting int8=True.")
self.args.int8 = True self.args.int8 = True
@ -369,6 +370,7 @@ class Exporter:
"batch": self.args.batch, "batch": self.args.batch,
"imgsz": self.imgsz, "imgsz": self.imgsz,
"names": model.names, "names": model.names,
"args": {k: v for k, v in self.args if k in fmt_keys},
} # model metadata } # model metadata
if model.task == "pose": if model.task == "pose":
self.metadata["kpt_shape"] = model.model[-1].kpt_shape self.metadata["kpt_shape"] = model.model[-1].kpt_shape