[Bug Fix] yolov8 export sample fix

[dsikka]

Summary

• There was a bug when running the following command:
  sparseml.ultralytics.export_onnx --model model.pt --export-samples 20 where the samples could not be exported.
• This is because, the export function pulls in the yaml path saved in the checkpointed model, which is local to wherever this mode was trained. To fix this, check if the yaml path (stored in args[data]) exists. If it does not, just use the default config provided when exporting the data samples.

Furthermore, there was one more check:

 if args["dataset_path"] is not None:
    args["data"] = data_from_dataset_path(
        args["data"], args["dataset_path"]
    )

• Here, if the user provides the dataset_path when running sparseml.ultralytics.export_onnx, the function data_from_dataset_path will attempt to pull the yaml from ultralytics. However, this will only work if the input to the function args["data"] stores the name of one of the yamls, not the path to a locally saved yaml, in which case it will try and downloading this yaml. Do we expect this case to exist?

Testing

• Tested locally using the command sparseml.ultralytics.export_onnx --model model.pt --export-samples 20

[bfineran]

LGTM - @anmarques can you test locally before landing

[anmarques]

Not really working on my environment.

Testing:

git clone git@git.neuralmagic.com:neuralmagic/yolov8-s-coco-pruned65.git
cd yolov8-s-coco-pruned65
sparseml.ultralytics.export_onnx --model training/model.pt --export-samples 20

Error:

Dataset 'coco128.yaml' not found ⚠️, missing paths ['/home/ubuntu/datasets/coco128/images/train2017']
Downloading https://ultralytics.com/assets/coco128.zip to /home/alexandre/test/coco128.zip...
100%|██████████████████████████████████████| 6.66M/6.66M [00:01<00:00, 6.69MB/s]
Unzipping /home/alexandre/test/coco128.zip to /home/alexandre/test...
Dataset download success ✅ (1.8s), saved to /home/alexandre/test

Traceback (most recent call last):
  File "/home/alexandre/environments/sparseml/lib/python3.8/site-packages/ultralytics/yolo/data/dataloaders/v5loader.py", line 482, in __init__
    raise FileNotFoundError(f'{prefix}{p} does not exist')
FileNotFoundError: /home/ubuntu/datasets/coco128/images/train2017 does not exist

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/home/alexandre/environments/sparseml/bin/sparseml.ultralytics.export_onnx", line 33, in <module>
    sys.exit(load_entry_point('sparseml-nightly', 'console_scripts', 'sparseml.ultralytics.export_onnx')())
  File "/home/alexandre/environments/sparseml/lib/python3.8/site-packages/click/core.py", line 1130, in __call__
    return self.main(*args, **kwargs)
  File "/home/alexandre/environments/sparseml/lib/python3.8/site-packages/click/core.py", line 1055, in main
    rv = self.invoke(ctx)
  File "/home/alexandre/environments/sparseml/lib/python3.8/site-packages/click/core.py", line 1404, in invoke
    return ctx.invoke(self.callback, **ctx.params)
  File "/home/alexandre/environments/sparseml/lib/python3.8/site-packages/click/core.py", line 760, in invoke
    return __callback(*args, **kwargs)
  File "/home/alexandre/git/sparseml/src/sparseml/yolov8/export.py", line 91, in main
    model.export(**kwargs)
  File "/home/alexandre/git/sparseml/src/sparseml/yolov8/trainers.py", line 778, in export
    data_loader, _ = create_dataloader(
  File "/home/alexandre/environments/sparseml/lib/python3.8/site-packages/ultralytics/yolo/data/dataloaders/v5loader.py", line 127, in create_dataloader
    dataset = LoadImagesAndLabels(
  File "/home/alexandre/environments/sparseml/lib/python3.8/site-packages/ultralytics/yolo/data/dataloaders/v5loader.py", line 487, in __init__
    raise FileNotFoundError(f'{prefix}Error loading data from {path}: {e}\n{HELP_URL}') from e
FileNotFoundError: Error loading data from /home/ubuntu/datasets/coco128/images/train2017: /home/ubuntu/datasets/coco128/images/train2017 does not exist
See https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data

It tried to use the default coco128 dataset. Since it didn't find it, it downloaded it. But later it tried to use in a different path (probably a path that is registered in config.yaml from using this dataset in a different occasion).

[anmarques]

Also, I think it's great to try to harden the export_onnx pathway, but we still have the issue that when we save the model it saves the local path to the dataset. In my view we should only save the name of the yaml file (such as 'coco.yaml') instead of the full path

[anmarques]

I verified the PR fixes the issue of the onnx export.

As I discussed w/ Dipika, we also need to fix it to not store the local dataset path when saving the checkpoint.