Browse Source

Fix `torch.hub.list('ultralytics/yolov5')` pathlib bug (#3921)

modifyDataloader
Glenn Jocher GitHub 3 years ago
parent
commit
411842e058
No known key found for this signature in database GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 5 additions and 5 deletions
  1. +5
    -5
      hubconf.py

+ 5
- 5
hubconf.py View File

@@ -4,12 +4,9 @@ Usage:
import torch
model = torch.hub.load('ultralytics/yolov5', 'yolov5s')
"""
from pathlib import Path

import torch

FILE = Path(__file__).absolute()


def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None):
"""Creates a specified YOLOv5 model
@@ -26,15 +23,18 @@ def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbo
Returns:
YOLOv5 pytorch model
"""
from pathlib import Path

from models.yolo import Model, attempt_load
from utils.general import check_requirements, set_logging
from utils.google_utils import attempt_download
from utils.torch_utils import select_device

check_requirements(requirements=FILE.parent / 'requirements.txt', exclude=('tensorboard', 'thop', 'opencv-python'))
file = Path(__file__).absolute()
check_requirements(requirements=file.parent / 'requirements.txt', exclude=('tensorboard', 'thop', 'opencv-python'))
set_logging(verbose=verbose)

save_dir = Path('') if str(name).endswith('.pt') else FILE.parent
save_dir = Path('') if str(name).endswith('.pt') else file.parent
path = (save_dir / name).with_suffix('.pt') # checkpoint path
try:
device = select_device(('0' if torch.cuda.is_available() else 'cpu') if device is None else device)

Loading…
Cancel
Save