ultralytics 8.0.206 engine Trainer updates (#6111)

Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: jamjamjon <51357717+jamjamjon@users.noreply.github.com>
This commit is contained in:
Glenn Jocher 2023-11-04 02:57:35 +01:00 committed by GitHub
parent 25bd3b9834
commit f2f5ed2c5e
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
7 changed files with 42 additions and 34 deletions

View file

@ -9,7 +9,6 @@ from ultralytics.cfg import TASK2DATA, get_cfg, get_save_dir
from ultralytics.hub.utils import HUB_WEB_ROOT
from ultralytics.nn.tasks import attempt_load_one_weight, guess_model_task, nn, yaml_model_load
from ultralytics.utils import ASSETS, DEFAULT_CFG_DICT, LOGGER, RANK, callbacks, checks, emojis, yaml_load
from ultralytics.utils.downloads import GITHUB_ASSETS_STEMS
class Model(nn.Module):
@ -88,10 +87,8 @@ class Model(nn.Module):
return
# Load or create new YOLO model
suffix = Path(model).suffix
if not suffix and Path(model).stem in GITHUB_ASSETS_STEMS:
model, suffix = Path(model).with_suffix('.pt'), '.pt' # add suffix, i.e. yolov8n -> yolov8n.pt
if suffix in ('.yaml', '.yml'):
model = checks.check_model_file_from_stem(model) # add suffix, i.e. yolov8n -> yolov8n.pt
if Path(model).suffix in ('.yaml', '.yml'):
self._new(model, task)
else:
self._load(model, task)