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

ultralytics 8.0.44 export and task fixes (#1088)

Browse source 
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Mehran Ghandehari <mehran.maps@gmail.com>
Co-authored-by: Laughing <61612323+Laughing-q@users.noreply.github.com>
This commit is contained in:
Glenn Jocher 2023-02-24 03:11:25 +01:00 • committed by GitHub
parent fe61018975
commit 3ea659411b
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
32 changed files with 439 additions and 480 deletions
Download patch file Download diff file Expand all files Collapse all files

139
ultralytics/yolo/engine/model.py
View file

@ -9,76 +9,72 @@ from ultralytics.nn.tasks import (ClassificationModel, DetectionModel, Segmentat
guess_model_task, nn)
from ultralytics.yolo.cfg import get_cfg
from ultralytics.yolo.engine.exporter import Exporter
from ultralytics.yolo.utils import DEFAULT_CFG, DEFAULT_CFG_DICT, LOGGER, RANK, callbacks, yaml_load
from ultralytics.yolo.utils import DEFAULT_CFG, DEFAULT_CFG_DICT, DEFAULT_CFG_KEYS, LOGGER, RANK, callbacks, yaml_load
from ultralytics.yolo.utils.checks import check_file, check_imgsz, check_yaml
from ultralytics.yolo.utils.downloads import GITHUB_ASSET_STEMS
from ultralytics.yolo.utils.torch_utils import smart_inference_mode
# Map head to model, trainer, validator, and predictor classes
MODEL_MAP = {
TASK_MAP = {
'classify': [
ClassificationModel, 'yolo.TYPE.classify.ClassificationTrainer', 'yolo.TYPE.classify.ClassificationValidator',
'yolo.TYPE.classify.ClassificationPredictor'],
ClassificationModel, yolo.v8.classify.ClassificationTrainer, yolo.v8.classify.ClassificationValidator,
yolo.v8.classify.ClassificationPredictor],
'detect': [
DetectionModel, 'yolo.TYPE.detect.DetectionTrainer', 'yolo.TYPE.detect.DetectionValidator',
'yolo.TYPE.detect.DetectionPredictor'],
DetectionModel, yolo.v8.detect.DetectionTrainer, yolo.v8.detect.DetectionValidator,
yolo.v8.detect.DetectionPredictor],
'segment': [
SegmentationModel, 'yolo.TYPE.segment.SegmentationTrainer', 'yolo.TYPE.segment.SegmentationValidator',
'yolo.TYPE.segment.SegmentationPredictor']}
SegmentationModel, yolo.v8.segment.SegmentationTrainer, yolo.v8.segment.SegmentationValidator,
yolo.v8.segment.SegmentationPredictor]}
class YOLO:
"""
YOLO (You Only Look Once) object detection model.
YOLO (You Only Look Once) object detection model.
Args:
model (str, Path): Path to the model file to load or create.
type (str): Type/version of models to use. Defaults to "v8".
Args:
model (str, Path): Path to the model file to load or create.
Attributes:
type (str): Type/version of models being used.
ModelClass (Any): Model class.
TrainerClass (Any): Trainer class.
ValidatorClass (Any): Validator class.
PredictorClass (Any): Predictor class.
predictor (Any): Predictor object.
model (Any): Model object.
trainer (Any): Trainer object.
task (str): Type of model task.
ckpt (Any): Checkpoint object if model loaded from *.pt file.
cfg (str): Model configuration if loaded from *.yaml file.
ckpt_path (str): Checkpoint file path.
overrides (dict): Overrides for trainer object.
metrics_data (Any): Data for metrics.
Attributes:
predictor (Any): The predictor object.
model (Any): The model object.
trainer (Any): The trainer object.
task (str): The type of model task.
ckpt (Any): The checkpoint object if the model loaded from *.pt file.
cfg (str): The model configuration if loaded from *.yaml file.
ckpt_path (str): The checkpoint file path.
overrides (dict): Overrides for the trainer object.
metrics_data (Any): The data for metrics.
Methods:
__call__(): Alias for predict method.
_new(cfg, verbose=True): Initializes a new model and infers the task type from the model definitions.
_load(weights): Initializes a new model and infers the task type from the model head.
_check_is_pytorch_model(): Raises TypeError if model is not a PyTorch model.
reset(): Resets the model modules.
info(verbose=False): Logs model info.
fuse(): Fuse model for faster inference.
predict(source=None, stream=False, **kwargs): Perform prediction using the YOLO model.
Methods:
__call__(source=None, stream=False, **kwargs):
Alias for the predict method.
_new(cfg:str, verbose:bool=True) -> None:
Initializes a new model and infers the task type from the model definitions.
_load(weights:str, task:str='') -> None:
Initializes a new model and infers the task type from the model head.
_check_is_pytorch_model() -> None:
Raises TypeError if the model is not a PyTorch model.
reset() -> None:
Resets the model modules.
info(verbose:bool=False) -> None:
Logs the model info.
fuse() -> None:
Fuses the model for faster inference.
predict(source=None, stream=False, **kwargs) -> List[ultralytics.yolo.engine.results.Results]:
Performs prediction using the YOLO model.
Returns:
list(ultralytics.yolo.engine.results.Results): The prediction results.
"""
Returns:
list[ultralytics.yolo.engine.results.Results]: The prediction results.
"""
def __init__(self, model='yolov8n.pt', type='v8') -> None:
def __init__(self, model='yolov8n.pt') -> None:
"""
Initializes the YOLO model.
Args:
model (str, Path): model to load or create
type (str): Type/version of models to use. Defaults to "v8".
"""
self._reset_callbacks()
self.type = type
self.ModelClass = None # model class
self.TrainerClass = None # trainer class
self.ValidatorClass = None # validator class
self.PredictorClass = None # predictor class
self.predictor = None # reuse predictor
self.model = None # model object
self.trainer = None # trainer object
@ -101,6 +97,10 @@ class YOLO:
def __call__(self, source=None, stream=False, **kwargs):
return self.predict(source, stream, **kwargs)
def __getattr__(self, attr):
name = self.__class__.__name__
raise AttributeError(f"'{name}' object has no attribute '{attr}'. See valid attributes below.\n{self.__doc__}")
def _new(self, cfg: str, verbose=True):
"""
Initializes a new model and infers the task type from the model definitions.
@ -112,11 +112,15 @@ class YOLO:
self.cfg = check_yaml(cfg) # check YAML
cfg_dict = yaml_load(self.cfg, append_filename=True) # model dict
self.task = guess_model_task(cfg_dict)
self.ModelClass, self.TrainerClass, self.ValidatorClass, self.PredictorClass = self._assign_ops_from_task()
self.model = self.ModelClass(cfg_dict, verbose=verbose and RANK == -1) # initialize
self.model = TASK_MAP[self.task][0](cfg_dict, verbose=verbose and RANK == -1) # build model
self.overrides['model'] = self.cfg
def _load(self, weights: str):
# Below added to allow export from yamls
args = {**DEFAULT_CFG_DICT, **self.overrides} # combine model and default args, preferring model args
self.model.args = {k: v for k, v in args.items() if k in DEFAULT_CFG_KEYS} # attach args to model
self.model.task = self.task
def _load(self, weights: str, task=''):
"""
Initializes a new model and infers the task type from the model head.
@ -127,8 +131,7 @@ class YOLO:
if suffix == '.pt':
self.model, self.ckpt = attempt_load_one_weight(weights)
self.task = self.model.args['task']
self.overrides = self.model.args
self._reset_ckpt_args(self.overrides)
self.overrides = self.model.args = self._reset_ckpt_args(self.model.args)
self.ckpt_path = self.model.pt_path
else:
weights = check_file(weights)
@ -136,7 +139,6 @@ class YOLO:
self.task = guess_model_task(weights)
self.ckpt_path = weights
self.overrides['model'] = weights
self.ModelClass, self.TrainerClass, self.ValidatorClass, self.PredictorClass = self._assign_ops_from_task()
def _check_is_pytorch_model(self):
"""
@ -189,12 +191,13 @@ class YOLO:
"""
overrides = self.overrides.copy()
overrides['conf'] = 0.25
overrides.update(kwargs)
overrides.update(kwargs) # prefer kwargs
overrides['mode'] = kwargs.get('mode', 'predict')
assert overrides['mode'] in ['track', 'predict']
overrides['save'] = kwargs.get('save', False) # not save files by default
if not self.predictor:
self.predictor = self.PredictorClass(overrides=overrides)
self.task = overrides.get('task') or self.task
self.predictor = TASK_MAP[self.task][3](overrides=overrides)
self.predictor.setup_model(model=self.model)
else: # only update args if predictor is already setup
self.predictor.args = get_cfg(self.predictor.args, overrides)
@ -226,12 +229,15 @@ class YOLO:
overrides['mode'] = 'val'
args = get_cfg(cfg=DEFAULT_CFG, overrides=overrides)
args.data = data or args.data
args.task = self.task
if 'task' in overrides:
self.task = args.task
else:
args.task = self.task
if args.imgsz == DEFAULT_CFG.imgsz and not isinstance(self.model, (str, Path)):
args.imgsz = self.model.args['imgsz'] # use trained imgsz unless custom value is passed
args.imgsz = check_imgsz(args.imgsz, max_dim=1)
validator = self.ValidatorClass(args=args)
validator = TASK_MAP[self.task][2](args=args)
validator(model=self.model)
self.metrics_data = validator.metrics
@ -267,8 +273,7 @@ class YOLO:
args.imgsz = self.model.args['imgsz'] # use trained imgsz unless custom value is passed
if args.batch == DEFAULT_CFG.batch:
args.batch = 1 # default to 1 if not modified
exporter = Exporter(overrides=args)
return exporter(model=self.model)
return Exporter(overrides=args)(model=self.model)
def train(self, **kwargs):
"""
@ -282,15 +287,15 @@ class YOLO:
overrides.update(kwargs)
if kwargs.get('cfg'):
LOGGER.info(f"cfg file passed. Overriding default params with {kwargs['cfg']}.")
overrides = yaml_load(check_yaml(kwargs['cfg']), append_filename=True)
overrides['task'] = self.task
overrides = yaml_load(check_yaml(kwargs['cfg']))
overrides['mode'] = 'train'
if not overrides.get('data'):
raise AttributeError("Dataset required but missing, i.e. pass 'data=coco128.yaml'")
if overrides.get('resume'):
overrides['resume'] = self.ckpt_path
self.trainer = self.TrainerClass(overrides=overrides)
self.task = overrides.get('task') or self.task
self.trainer = TASK_MAP[self.task][1](overrides=overrides)
if not overrides.get('resume'): # manually set model only if not resuming
self.trainer.model = self.trainer.get_model(weights=self.model if self.ckpt else None, cfg=self.model.yaml)
self.model = self.trainer.model
@ -311,13 +316,6 @@ class YOLO:
self._check_is_pytorch_model()
self.model.to(device)
def _assign_ops_from_task(self):
model_class, train_lit, val_lit, pred_lit = MODEL_MAP[self.task]
trainer_class = eval(train_lit.replace('TYPE', f'{self.type}'))
validator_class = eval(val_lit.replace('TYPE', f'{self.type}'))
predictor_class = eval(pred_lit.replace('TYPE', f'{self.type}'))
return model_class, trainer_class, validator_class, predictor_class
@property
def names(self):
"""
@ -357,9 +355,8 @@ class YOLO:
@staticmethod
def _reset_ckpt_args(args):
for arg in 'augment', 'verbose', 'project', 'name', 'exist_ok', 'resume', 'batch', 'epochs', 'cache', \
'save_json', 'half', 'v5loader', 'device', 'cfg', 'save', 'rect', 'plots', 'opset', 'simplify':
args.pop(arg, None)
include = {'imgsz', 'data', 'task', 'single_cls'} # only remember these arguments when loading a PyTorch model
return {k: v for k, v in args.items() if k in include}
@staticmethod
def _reset_callbacks():