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:
parent
25bd3b9834
commit
f2f5ed2c5e
7 changed files with 42 additions and 34 deletions
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue