Issue Saving MegaDetectorV6: Cannot Pickle '_thread.lock' Object

[nee1k]

When saving the MegaDetectorV6 model using the following code:

import torch
from PytorchWildlife.models import detection as pw_detection

model = pw_detection.MegaDetectorV6(
    device="cpu",
    pretrained=True,
    version="MDV6-yolov10-e"
)

torch.save(model, "MegaDetectorV6.pt")

I encountered the error:

TypeError: cannot pickle '_thread.lock' object

Could you please advise on: A recommended method to export this model directly to Hugging Face Transformers? How to save the model as a .pt file or convert it to ONNX given this serialization issue?

[danielaruizl1]

Hi @nee1k! It seems that the current issue with saving the model is due to the use of torch.save(). We recommend saving the model’s state dictionary instead. Please try the following approach:

import torch
from PytorchWildlife.models import detection as pw_detection

model = pw_detection.MegaDetectorV6(
    device="cpu",
    pretrained=True,
    version="MDV6-yolov10-e"
)

# Save state dictionary
torch.save(model.state_dict(), "MegaDetectorV6.pt")

Additionally, we haven’t tested compatibility with the HuggingFace Transformers package yet. However, our MegaDetectorv6 weights are based on unmodified YOLO architectures. As long as the HuggingFace Transformers library can load the detection model, you should be able to load the state dictionary without any issues. Here’s a brief example of how to load the state dict:

# Load state dictionary
model.load_state_dict(torch.load("MegaDetectorV6.pt"))

MegaDetectorV6 is trained with 3 classes (animal, person, vehicle), while most YOLO models pretrained on COCO have 80 classes. Make sure to modify the detection head to reflect this difference.

Please let us know if this approach works!