sparseml icon indicating copy to clipboard operation
sparseml copied to clipboard

[SparseZoo v2 Bridge] Save yolo training artifacts to ModelDirectory directory

Open dbogunowicz opened this issue 3 years ago • 0 comments

Allows creating the ModelDirectory local directory from the direct output of the yolo training procedure.

To test:

  1. You need to install locally the sparsezoo's feature/ci-cd-refactor/validations branch
  2. Then run:
import os.path
from src.sparseml.yolov5.helpers import get_model_directory
from sparsezoo.v2.model_directory import ModelDirectory

# where the training artifacts are
path_to_training_outputs = "/home/damian/yolo/yolov5_runs/train/exp5"
model_file_torch = "last.pt"
save_dir = "/home/damian/my_model_dir"

# create model card if none exists
model_card_path = os.path.join(path_to_training_outputs, "model.md")
if not os.path.exists(model_card_path):
    open(model_card_path, "w").close()


get_model_directory(
    output_dir=save_dir,
    training_outputs_dir=path_to_training_outputs,
    model_file_torch = model_file_torch
)

# print the resulting directory
start_path = save_dir
for path, dirs, files in os.walk(start_path):
    for filename in files:
        print(os.path.join(path, filename))

# this will fail in a controlled manner (the model card is empty)
assert ModelDirectory.from_directory(save_dir)

Output:

/home/damian/my_model_dir/last.onnx
/home/damian/my_model_dir/model.md
/home/damian/my_model_dir/sample_outputs.tar.gz
/home/damian/my_model_dir/sample_inputs.tar.gz
/home/damian/my_model_dir/training/last.pt
/home/damian/my_model_dir/deployment/last.onnx
/home/damian/my_model_dir/logs/events.out.tfevents.1656931633.lambdaquad.2383394.0

Traceback (most recent call last):
  File "/home/damian/sparseml/src/sparseml/yolov5/test.py", line 29, in <module>
    assert ModelDirectory.from_directory(save_dir)
  File "/home/damian/sparsezoo/src/sparsezoo/v2/model_directory.py", line 196, in from_directory
    return ModelDirectory(files=files, name="model_directory", path=directory_path)
  File "/home/damian/sparsezoo/src/sparsezoo/v2/model_directory.py", line 84, in __init__
    ] = self._sample_outputs_list_to_dict(
  File "/home/damian/sparsezoo/src/sparsezoo/v2/model_directory.py", line 474, in _sample_outputs_list_to_dict
    framework_name = self.model_card._validate_model_card()["framework"]
TypeError: 'NoneType' object is not subscriptable

dbogunowicz avatar Jul 04 '22 09:07 dbogunowicz