diff --git a/ultralytics/yolo/engine/trainer.py b/ultralytics/yolo/engine/trainer.py index 44decbd..4bf6947 100644 --- a/ultralytics/yolo/engine/trainer.py +++ b/ultralytics/yolo/engine/trainer.py @@ -24,13 +24,12 @@ from ultralytics.yolo.utils import LOGGER, ROOT from ultralytics.yolo.utils.files import increment_path, save_yaml from ultralytics.yolo.utils.modeling import get_model -CONFIG_PATH_ABS = ROOT / "yolo/utils/configs" -DEFAULT_CONFIG = "defaults.yaml" +DEFAULT_CONFIG = ROOT / "yolo/utils/configs/default.yml" class BaseTrainer: - def __init__(self, config=CONFIG_PATH_ABS / DEFAULT_CONFIG, overrides={}): + def __init__(self, config=DEFAULT_CONFIG, overrides={}): self.console = LOGGER self.args = self._get_config(config, overrides) self.validator = None diff --git a/ultralytics/yolo/utils/configs/defaults.yaml b/ultralytics/yolo/utils/configs/default.yml similarity index 65% rename from ultralytics/yolo/utils/configs/defaults.yaml rename to ultralytics/yolo/utils/configs/default.yml index b2e7474..8508744 100644 --- a/ultralytics/yolo/utils/configs/defaults.yaml +++ b/ultralytics/yolo/utils/configs/default.yml @@ -1,25 +1,27 @@ -model: null -data: null +# YOLO 🚀 by Ultralytics, GPL-3.0 license +# Default training settings and hyperparameters for medium-augmentation COCO training -# Training options + +# Train settings ------------------------------------------------------------------------------------------------------- +model: null # i.e. yolov5s.pt +data: null # i.e. coco128.yaml epochs: 300 batch_size: 16 img_size: 640 nosave: False -cache: False # True/ram for ram, or disc -device: '' # cuda device, i.e. 0 or 0,1,2,3 or cpu +cache: False # True/ram, disk or False +device: '' # cuda device, i.e. 0 or 0,1,2,3 or cpu workers: 8 -project: "ultralytics-yolo" -name: "exp" # TODO: make this informative, maybe exp{#number}_{datetime} ? +project: 'runs' +name: 'exp' exist_ok: False pretrained: False -optimizer: "Adam" # choices=['SGD', 'Adam', 'AdamW', 'RMSProp'] +optimizer: 'SGD' # choices=['SGD', 'Adam', 'AdamW', 'RMSProp'] verbose: False seed: 0 local_rank: -1 -#-----------------------------------# -# Hyper-parameters +# Hyperparameters ------------------------------------------------------------------------------------------------------ lr0: 0.001 # initial learning rate (SGD=1E-2, Adam=1E-3) lrf: 0.01 # final OneCycleLR learning rate (lr0 * lrf) momentum: 0.937 # SGD momentum/Adam beta1 @@ -50,9 +52,8 @@ mosaic: 1.0 # image mosaic (probability) mixup: 0.0 # image mixup (probability) copy_paste: 0.0 # segment copy-paste (probability) -# Hydra configs ------------------------------------- -# to disable hydra directory creation +# Hydra configs -------------------------------------------------------------------------------------------------------- hydra: - output_subdir: null + output_subdir: null # disable hydra directory creation run: dir: . diff --git a/ultralytics/yolo/utils/modeling/__init__.py b/ultralytics/yolo/utils/modeling/__init__.py index c3328c7..35030bb 100644 --- a/ultralytics/yolo/utils/modeling/__init__.py +++ b/ultralytics/yolo/utils/modeling/__init__.py @@ -107,18 +107,17 @@ def parse_model(d, ch): # model_dict, input_channels(3) return nn.Sequential(*layers), sorted(save) -def get_model(model: str): +def get_model(model='s.pt', pretrained=True): + # Load a YOLO model locally, from torchvision, or from Ultralytics assets if model.endswith(".pt"): model = model.split(".")[0] - if Path(model + ".pt").is_file(): - trained_model = torch.load(model + ".pt", map_location='cpu') - elif model in torchvision.models.__dict__: # try torch hub classifier models - trained_model = torch.hub.load("pytorch/vision", model, pretrained=True) - else: - model_ckpt = attempt_download(model + ".pt") # try ultralytics assets - trained_model = torch.load(model_ckpt, map_location='cpu') - return trained_model + if Path(f"{model}.pt").is_file(): # local file + return torch.load(f"{model}.pt", map_location='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') def yaml_load(file='data.yaml'): diff --git a/ultralytics/yolo/v8/classify/train.py b/ultralytics/yolo/v8/classify/train.py index 6cc21bb..89c7041 100644 --- a/ultralytics/yolo/v8/classify/train.py +++ b/ultralytics/yolo/v8/classify/train.py @@ -4,13 +4,13 @@ from pathlib import Path import hydra import torch -import torchvision from ultralytics.yolo import v8 from ultralytics.yolo.data import build_classification_dataloader -from ultralytics.yolo.engine.trainer import CONFIG_PATH_ABS, DEFAULT_CONFIG, BaseTrainer +from ultralytics.yolo.engine.trainer import DEFAULT_CONFIG, BaseTrainer from ultralytics.yolo.utils.downloads import download from ultralytics.yolo.utils.files import WorkingDirectory +from ultralytics.yolo.utils.loggers import colorstr from ultralytics.yolo.utils.torch_utils import LOCAL_RANK, torch_distributed_zero_first @@ -30,8 +30,7 @@ class ClassificationTrainer(BaseTrainer): else: url = f'https://github.com/ultralytics/yolov5/releases/download/v1.0/{dataset}.zip' download(url, dir=data_dir.parent) - # TODO: add colorstr - s = f"Dataset download success ✅ ({time.time() - t:.1f}s), saved to {'bold', data_dir}\n" + s = f"Dataset download success ✅ ({time.time() - t:.1f}s), saved to {colorstr('bold', data_dir)}\n" self.console.info(s) train_set = data_dir / "train" test_set = data_dir / 'test' if (data_dir / 'test').exists() else data_dir / 'val' # data/test or data/val @@ -48,7 +47,7 @@ class ClassificationTrainer(BaseTrainer): return torch.nn.functional.cross_entropy(preds, targets) -@hydra.main(version_base=None, config_path=CONFIG_PATH_ABS, config_name=str(DEFAULT_CONFIG).split(".")[0]) +@hydra.main(version_base=None, config_path=DEFAULT_CONFIG.parent, config_name=DEFAULT_CONFIG.stem) def train(cfg): cfg.model = cfg.model or "resnet18" cfg.data = cfg.data or "imagenette160" # or yolo.ClassificationDataset("mnist") @@ -59,7 +58,7 @@ def train(cfg): if __name__ == "__main__": """ CLI usage: - python ../path/to/train.py args.epochs=10 args.project="name" hyps.lr0=0.1 + python path/to/train.py epochs=10 project=PROJECT lr0=0.1 TODO: Direct cli support, i.e, yolov8 classify_train args.epochs 10