Support both *.yml and *.yaml files (#4086)

Co-authored-by: ChristopherRogers1991 <ChristopherRogers1991@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Sungjoo(Dennis) Hwang <48212469+Denny-Hwang@users.noreply.github.com>

ultralytics/models/fastsam/model.py

@@ -23,7 +23,7 @@ class FastSAM(Model):
         """Call the __init__ method of the parent class (YOLO) with the updated default model"""
         if model == 'FastSAM.pt':
             model = 'FastSAM-x.pt'
-        assert Path(model).suffix != '.yaml', 'FastSAM models only support pre-trained models.'
+        assert Path(model).suffix not in ('.yaml', '.yml'), 'FastSAM models only support pre-trained models.'
         super().__init__(model=model, task='segment')
 
     @property

ultralytics/models/nas/model.py

@@ -23,7 +23,7 @@ from .val import NASValidator
 class NAS(Model):
 
     def __init__(self, model='yolo_nas_s.pt') -> None:
-        assert Path(model).suffix != '.yaml', 'YOLO-NAS models only support pre-trained models.'
+        assert Path(model).suffix not in ('.yaml', '.yml'), 'YOLO-NAS models only support pre-trained models.'
         super().__init__(model, task='detect')
 
     @smart_inference_mode()

ultralytics/models/rtdetr/model.py

@@ -16,8 +16,8 @@ class RTDETR(Model):
     """
 
     def __init__(self, model='rtdetr-l.pt') -> None:
-        if model and not model.endswith('.pt') and not model.endswith('.yaml'):
-            raise NotImplementedError('RT-DETR only supports creating from pt file or yaml file.')
+        if model and not model.split('.')[-1] in ('pt', 'yaml', 'yml'):
+            raise NotImplementedError('RT-DETR only supports creating from *.pt file or *.yaml file.')
         super().__init__(model=model, task='detect')
 
     @property

ultralytics/models/yolo/classify/train.py

@@ -57,7 +57,7 @@ class ClassificationTrainer(BaseTrainer):
             self.model, _ = attempt_load_one_weight(model, device='cpu')
             for p in self.model.parameters():
                 p.requires_grad = True  # for training
-        elif model.endswith('.yaml'):
+        elif model.split('.')[-1] in ('yaml', 'yml'):
             self.model = self.get_model(cfg=model)
         elif model in torchvision.models.__dict__:
             self.model = torchvision.models.__dict__[model](weights='IMAGENET1K_V1' if self.args.pretrained else None)