Remap names in a YOLOv5 PyTorch model file

[kevinsbarnard]

Requirements

• Python
• PyTorch: pip install torch
• A YOLOv5 saved model file, e.g.: model.pt

Steps

1. Load the YOLOv5 PyTorch model file. This results in a Python dict in which the actual YOLOv5 model is stored in the model field.

import torch

model_dict = torch.load("/path/to/model.pt")
model = model_dict["model"]

2. Change the names. The model.names field is a Python list that contains the names. Make sure you do not change the ordering of the names list. For example:

name_index = model.names.index("old-name")
model.names[name_index] = "new-name"

3. Save the updated model dict.

with open("/path/to/updated_model.pt", "wb") as f:
    torch.save(model, f)

4. (optional) Confirm the old name is no longer in the model (and the new name is) with grep:

$ grep 'old-name' /path/to/updated_model.pt
$ grep 'new-name' /path/to/updated_model.pt
grep: updated_model.pt: binary file matches