HongChuang/ultralytics-ascend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 3

standalone val (#56)

Browse source 
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Ayush Chaurasia 2022-11-30 15:04:44 +05:30 • committed by GitHub
parent 3a241e4cea
commit 5a52e7663a
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
16 changed files with 161 additions and 31 deletions
Download patch file Download diff file Expand all files Collapse all files

4
ultralytics/yolo/utils/modeling/__init__.py
View file

@ -113,8 +113,8 @@ def get_model(model='s.pt', pretrained=True):
model = model.split(".")[0]
if Path(f"{model}.pt").is_file(): # local file
return torch.load(f"{model}.pt", map_location='cpu')
return attempt_load_weights(f"{model}.pt", device='cpu')
elif model in torchvision.models.__dict__: # TorchVision models i.e. resnet50, efficientnet_b0
return torchvision.models.__dict__[model](weights='IMAGENET1K_V1' if pretrained else None)
else: # Ultralytics assets
return torch.load(attempt_download(f"{model}.pt"), map_location='cpu')
return attempt_load_weights(f"{model}.pt", device='cpu')