Skip to content

Commit

Permalink
Fix loading Yolov8 sparsezoo models (#1561)
Browse files Browse the repository at this point in the history
* Fix loading Yolov8 sparsezoo models

* quality checks
  • Loading branch information
abhinavnmagic committed May 17, 2023
1 parent b96a89a commit 278c24a
Show file tree
Hide file tree
Showing 2 changed files with 3 additions and 0 deletions.
2 changes: 2 additions & 0 deletions src/sparseml/yolov8/trainers.py
Original file line number Diff line number Diff line change
Expand Up @@ -534,6 +534,7 @@ def __init__(self, model="yolov8n.yaml", type="v8") -> None:
model = download_framework_model_by_recipe_type(
Model(model_str), model_suffix="pt"
)
model_str = str(model)
self.is_sparseml_checkpoint = True

if model_str.endswith(".pt"):
Expand Down Expand Up @@ -800,6 +801,7 @@ def val(self, data=None, **kwargs):
overrides["rect"] = True # rect batches as default
overrides.update(kwargs)
overrides["mode"] = "val"
overrides["data"] = data or overrides["data"]
args = get_cfg(cfg=DEFAULT_CFG, overrides=overrides)
args.data = data or args.data
args.task = self.task
Expand Down
1 change: 1 addition & 0 deletions src/sparseml/yolov8/val.py
Original file line number Diff line number Diff line change
Expand Up @@ -76,6 +76,7 @@
default=None,
help="Path to override default datasets path.",
)
@click.option("--batch", default=16, type=int, help="number of images per batch")
def main(**kwargs):
if kwargs["dataset_path"] is not None:
kwargs["data"] = data_from_dataset_path(kwargs["data"], kwargs["dataset_path"])
Expand Down

0 comments on commit 278c24a

Please sign in to comment.