Skip to content

Reference for ultralytics/models/nas/model.py

Note

Full source code for this file is available at https://github.com/ultralytics/ultralytics/blob/main/ultralytics/models/nas/model.py. Help us fix any issues you see by submitting a Pull Request 🛠️. Thank you 🙏!


ultralytics.models.nas.model.NAS

Bases: Model

Source code in ultralytics/models/nas/model.py
class NAS(Model):

    def __init__(self, model='yolo_nas_s.pt') -> None:
        assert Path(model).suffix not in ('.yaml', '.yml'), 'YOLO-NAS models only support pre-trained models.'
        super().__init__(model, task='detect')

    @smart_inference_mode()
    def _load(self, weights: str, task: str):
        # Load or create new NAS model
        import super_gradients
        suffix = Path(weights).suffix
        if suffix == '.pt':
            self.model = torch.load(weights)
        elif suffix == '':
            self.model = super_gradients.training.models.get(weights, pretrained_weights='coco')
        # Standardize model
        self.model.fuse = lambda verbose=True: self.model
        self.model.stride = torch.tensor([32])
        self.model.names = dict(enumerate(self.model._class_names))
        self.model.is_fused = lambda: False  # for info()
        self.model.yaml = {}  # for info()
        self.model.pt_path = weights  # for export()
        self.model.task = 'detect'  # for export()

    def info(self, detailed=False, verbose=True):
        """
        Logs model info.

        Args:
            detailed (bool): Show detailed information about model.
            verbose (bool): Controls verbosity.
        """
        return model_info(self.model, detailed=detailed, verbose=verbose, imgsz=640)

    @property
    def task_map(self):
        return {'detect': {'predictor': NASPredictor, 'validator': NASValidator}}

info(detailed=False, verbose=True)

Logs model info.

Parameters:

Name Type Description Default
detailed bool

Show detailed information about model.

False
verbose bool

Controls verbosity.

True
Source code in ultralytics/models/nas/model.py
def info(self, detailed=False, verbose=True):
    """
    Logs model info.

    Args:
        detailed (bool): Show detailed information about model.
        verbose (bool): Controls verbosity.
    """
    return model_info(self.model, detailed=detailed, verbose=verbose, imgsz=640)




Created 2023-07-16, Updated 2023-08-07
Authors: glenn-jocher (5), Laughing-q (1)