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.

Sign up for GitHub

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

Jump to bottom

YOLOv8 model metadata on export #2353

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

YOLOv8 model metadata on export #2353

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 24 additions & 1 deletion src/sparseml/yolov8/trainers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@
from typing import List, Optional

import torch
import yaml

from sparseml.optim.helpers import load_recipe_yaml_str
from sparseml.pytorch.optim.manager import ScheduledModifierManager
Expand All @@ -46,7 +47,7 @@
SparseSegmentationValidator,
)
from sparsezoo import Model
from sparsezoo.utils import validate_onnx
from sparsezoo.utils import validate_onnx, load_model, save_onnx
from ultralytics import __version__
from ultralytics.nn.modules import Detect, Segment
from ultralytics.nn.tasks import SegmentationModel, attempt_load_one_weight
Expand Down Expand Up @@ -784,6 +785,28 @@ def export(self, **kwargs):
except Exception:
pass

# Add model metadata
data_yaml = Path(self.ckpt.get("train_args", {}).get("data", "")).read_text("utf-8")
names = yaml.safe_load(data_yaml).get("names", {})
metadata = {
"description": f"Ultralytics YOLOv8 via {source}",
"author": "Ultralytics",
"date": self.ckpt.get("date"),
"version": self.ckpt.get("version"),
"license": "AGPL-3.0 License (https://ultralytics.com/license)",
"docs": "https://docs.ultralytics.com",
"stride": int(max(model.model.stride)),
"task": self.task,
"batch": self.overrides.get("batch"),
"imgsz": self.overrides.get("imgsz"),
"names": names.copy(),
}
_onnx_m = load_model(complete_path)
for k,v in metadata.items():
meta = _onnx_m.metadata_props.add()
meta.key, meta.value = k, str(v)
_ = save_onnx(model=_onnx_m, model_path=complete_path)

validate_onnx(complete_path)
deployment_folder = exporter.create_deployment_folder(onnx_model_name=name)
if args["export_samples"]:
Expand Down