Saving Model in YOLOv8 format

[menonajayki]

Search before asking

• [X] I have searched the Ultralytics YOLO issues and discussions and found no similar questions.

Question

Hello Team,

I am currently working on federated learning with YOLOv8 and need to save the aggregated models in a format compatible with YOLOv8. I am using the following method to save the model:

def save_aggregated_model(parameters: List[np.ndarray], round_num: int):
    model = load_model()  # This loads a pretrained yolov8n.pt model
    set_parameters(model, parameters)  # This is part of flwr framework
    os.makedirs('model', exist_ok=True)
    model_save_path = f'model/FL_aggregated_round_{round_num}.pt'
    torch.save(model.state_dict(), model_save_path)
    logger.info(f"Model saved to {model_save_path} after round {round_num}")

def set_parameters(model, parameters: List[np.ndarray]):
    model_params = list(model.parameters())
    if len(model_params) != len(parameters):
        raise ValueError(f"Parameter count mismatch: model has {len(model_params)}, but received {len(parameters)}")
    for param, param_data in zip(model_params, parameters):
        param_ = torch.tensor(param_data)
        if param.data.shape != param_.data.shape:
            raise ValueError(f"Shape mismatch: model expects {param.data.shape}, but received {param_.data.shape}")
        param.data.copy_(param_.data)

While this saves the model's parameters, the saved checkpoints are in the format: _odictkeys(['model.model.0.conv.weight', 'model.model.0.bn.weight', 'model.model.0.bn.bias', .....], which is not directly usable for inference with YOLOv8. (odict.txt)

I would like to convert this to yolov8 Detection Model like this: _dict_keys(['date', 'version', 'license', 'docs', 'epoch', 'best_fitness', 'model', 'ema', 'updates', 'optimizer', 'train_args', 'train_metrics', 'trainresults']). Where the 'model' check point have the dictionary.

Could you provide guidance on how to adapt my current method to save the model in this format? Specifically, I need to include all the necessary components and metadata required by YOLOv8 for proper inference. Any suggestions or code examples on how to structure the saved model checkpoint in this format would be greatly appreciated.

Thank you!

Additional

No response

[github-actions[bot]]

👋 Hello @menonajayki, thank you for your interest in Ultralytics YOLOv8 🚀! We recommend a visit to the Docs for new users where you can find many Python and CLI usage examples and where many of the most common questions may already be answered.

If this is a 🐛 Bug Report, please provide a minimum reproducible example to help us debug it.

If this is a custom training ❓ Question, please provide as much information as possible, including dataset image examples and training logs, and verify you are following our Tips for Best Training Results.

Join the vibrant Ultralytics Discord 🎧 community for real-time conversations and collaborations. This platform offers a perfect space to inquire, showcase your work, and connect with fellow Ultralytics users.

Install

Pip install the ultralytics package including all requirements in a Python>=3.8 environment with PyTorch>=1.8.

pip install ultralytics

Environments

YOLOv8 may be run in any of the following up-to-date verified environments (with all dependencies including CUDA/CUDNN, Python and PyTorch preinstalled):

• Notebooks with free GPU:
• Google Cloud Deep Learning VM. See GCP Quickstart Guide
• Amazon Deep Learning AMI. See AWS Quickstart Guide
• Docker Image. See Docker Quickstart Guide

Status

If this badge is green, all Ultralytics CI tests are currently passing. CI tests verify correct operation of all YOLOv8 Modes and Tasks on macOS, Windows, and Ubuntu every 24 hours and on every commit.

[pderrenger]

Hello,

Thank you for reaching out. To save the aggregated model in a format compatible with YOLOv8, you need to ensure that the saved checkpoint includes the necessary metadata and structure expected by YOLOv8. The current method saves only the model parameters, but YOLOv8 checkpoints also include additional information such as training arguments, metrics, and optimizer state.

To achieve this, you can modify your saving function to include these components. Ensure that the saved dictionary contains keys like 'model', 'optimizer', 'train_args', etc., similar to the structure you mentioned.

If you encounter any issues or need further assistance, please ensure you are using the latest version of the Ultralytics package, as updates may include important fixes and improvements.

[menonajayki]

Hello @pderrenger,

Yes I understand that. I tried the following code to save the model,

# Load model saved to checkpoint
checkpoint1 = torch.load('yolov8n.pt')
checkpoint = torch.load('model/FL_aggregated_round_1.pt')

# Save the updated model.pt with new parameters
torch.save({
    'date': checkpoint1.get('date', ''),
    'version': checkpoint1.get('version', ''),
    'license': checkpoint1.get('license', ''),
    'docs': checkpoint1.get('docs', ''),
    'epoch': checkpoint1.get('epoch', 0),
    'best_fitness': checkpoint1.get('best_fitness', None),
    'model': checkpoint,
    'ema': checkpoint1.get('ema', None),
    'updates': checkpoint1.get('updates', None),
    'optimizer': checkpoint1.get('optimizer', None),
    'train_args': checkpoint1.get('train_args', {}),
    'train_metrics': checkpoint1.get('train_metrics', {}),
    'train_results': checkpoint1.get('train_results', {}),
}, 'updated_model1.pt')

But the model saved like this cannot be loaded in YOLO. When inspecting the model have checkpoints like this,

date: <class 'str'> with size: N/A version: <class 'str'> with size: N/A license: <class 'str'> with size: N/A docs: <class 'str'> with size: N/A epoch: <class 'int'> with size: N/A best_fitness: <class 'NoneType'> with size: N/A model: <class 'collections.OrderedDict'> with size: N/A ema: <class 'NoneType'> with size: N/A updates: <class 'NoneType'> with size: N/A optimizer: <class 'NoneType'> with size: N/A train_args: <class 'dict'> with size: N/A train_metrics: <class 'dict'> with size: N/A train_results: <class 'dict'> with size: N/A

And, when trying to load the model the YOLO is throwing error,

  File "C:\Users\AM\PycharmProjects\Imp Federated Learning\.venv\Lib\site-packages\ultralytics\models\yolo\model.py", line 23, in __init__
    super().__init__(model=model, task=task, verbose=verbose)
  File "C:\Users\AM\PycharmProjects\Imp Federated Learning\.venv\Lib\site-packages\ultralytics\engine\model.py", line 149, in __init__
    self._load(model, task=task)
  File "C:\Users\AM\PycharmProjects\Imp Federated Learning\.venv\Lib\site-packages\ultralytics\engine\model.py", line 230, in _load
    self.model, self.ckpt = attempt_load_one_weight(weights)
                            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\AM\PycharmProjects\Imp Federated Learning\.venv\Lib\site-packages\ultralytics\nn\tasks.py", line 857, in attempt_load_one_weight
    model = (ckpt.get("ema") or ckpt["model"]).to(device).float()  # FP32 model
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AttributeError: 'collections.OrderedDict' object has no attribute 'to'

I believe the issue is with model keypoint being ordered dictionary. How to save the keypoint 'model' in yolo format like

model: <class 'ultralytics.nn.tasks.DetectionModel'>

A sample code on achieving this will be really helpful.

[pderrenger]

Hello @menonajayki,

It looks like the issue arises because the 'model' key in your saved checkpoint is an OrderedDict rather than an instance of ultralytics.nn.tasks.DetectionModel. To resolve this, you need to ensure that the 'model' key contains the actual model object, not just its state dictionary.

Instead of directly saving the state dictionary under the 'model' key, you should first load the state dictionary into a YOLO model instance and then save the entire model object. This way, the 'model' key will correctly reference a DetectionModel instance.

Please ensure you are using the latest version of the Ultralytics package to avoid any compatibility issues. If you need further assistance, feel free to ask.

[Y-T-G]

model.ckpt["model"] = model.model
del model.ckpt["ema"]
model.save("saved.pt")

[menonajayki]

Hello @Y-T-G, I am using yolov8n.pt as reference and trying to load parameters from FL_aggregated_round_1.pt. But here the issues is that, FL_aggregated_round_1.pt is ordered dictionary and not ultralytics.nn.tasks.DetectionModel. I tried importing DetectionModel() but still the 'model' keypoint remains same since I am not able to assign an instance on DetectionModel to that dictionary. Do you have any sample implementation that I can look into, maybe I am missing something. I really appreciate your help!

[Y-T-G]

What do you keys look like?

[menonajayki]

@Y-T-G : odict.txt This is the keys of FL_aggregated_round_2.pt, I want to convert this to yolo format for inference. If you have some inputs this will be really helpful.