sparseml
sparseml copied to clipboard
[SparseZoo v2 Bridge] Save yolo training artifacts to ModelDirectory directory
Allows creating the ModelDirectory local directory from the direct output of the yolo training procedure.
To test:
- You need to install locally the sparsezoo's
feature/ci-cd-refactor/validationsbranch - 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