microsoft / CameraTraps

PyTorch Wildlife: a Collaborative Deep Learning Framework for Conservation.
https://cameratraps.readthedocs.io/en/latest/
MIT License
781 stars 246 forks source link

Problem loading detection model #489

Closed davidwhealey closed 6 months ago

davidwhealey commented 6 months ago

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?

davidwhealey commented 6 months ago

Found it. I had another module called models in my python path.