[X] I have searched the Pytorch-Wildlife issues and found no similar bug report.
Bug
I encounter a ModuleNotFoundError when trying to load the MegadetectorV5. The weights download, but I can't load the model. I get a similar error when loading by specifying the weights path.
Search before asking
Bug
I encounter a ModuleNotFoundError when trying to load the MegadetectorV5. The weights download, but I can't load the model. I get a similar error when loading by specifying the weights path.
ModuleNotFoundError Traceback (most recent call last) Cell In[9], line 2 1 device = 'cuda' if torch.cuda.is_available() else 'cpu' ----> 2 detection_model = pw_detection.MegaDetectorV5(device=device, pretrained=True)
File ~/miniconda3/envs/pytorch-wildlife-test/lib/python3.9/site-packages/PytorchWildlife/models/detection/yolov5/megadetector.py:44, in MegaDetectorV5.init(self, weights, device, pretrained) 41 else: 42 url = None ---> 44 super(MegaDetectorV5, self).init(weights=weights, device=device, url=url)
File ~/miniconda3/envs/pytorch-wildlife-test/lib/python3.9/site-packages/PytorchWildlife/models/detection/yolov5/base_detector.py:41, in YOLOV5Base.init(self, weights, device, url) 39 self.model = None 40 self.device = device ---> 41 self._load_model(weights, self.device, url) 42 self.model.to(self.device)
File ~/miniconda3/envs/pytorch-wildlife-test/lib/python3.9/site-packages/PytorchWildlife/models/detection/yolov5/base_detector.py:61, in YOLOV5Base._load_model(self, weights, device, url) 59 checkpoint = torch.load(weights, map_location=torch.device(device)) 60 elif url: ---> 61 checkpoint = load_state_dict_from_url(url, map_location=torch.device(self.device)) 62 else: 63 raise Exception("Need weights for inference.")
File ~/miniconda3/envs/pytorch-wildlife-test/lib/python3.9/site-packages/torch/hub.py:590, in load_state_dict_from_url(url, model_dir, map_location, progress, check_hash, file_name) ... 873 def find_class(self, mod_name, name): 874 mod_name = load_module_mapping.get(mod_name, mod_name) --> 875 return super().find_class(mod_name, name)
ModuleNotFoundError: No module named 'models.yolo'
Environment
Minimal Reproducible Example
import torch from PytorchWildlife.models import detection as pw_detection
device = 'cuda' if torch.cuda.is_available() else 'cpu' detection_model = pw_detection.MegaDetectorV5(device=device, pretrained=True)
Additional
No response
Are you willing to submit a PR?