From 29826241a00cff032e0739bad49c28bede59497f Mon Sep 17 00:00:00 2001 From: Void <94819873+WYYAHYT@users.noreply.github.com> Date: Fri, 29 Nov 2024 05:51:33 +0800 Subject: [PATCH] `__getattr__` support to access YOLO attributes via Model class (#17805) Co-authored-by: UltralyticsAssistant Co-authored-by: Ultralytics Assistant <135830346+UltralyticsAssistant@users.noreply.github.com> Co-authored-by: Glenn Jocher --- ultralytics/engine/model.py | 29 +++++++++++++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/ultralytics/engine/model.py b/ultralytics/engine/model.py index 667b54eb..874613d2 100644 --- a/ultralytics/engine/model.py +++ b/ultralytics/engine/model.py @@ -144,6 +144,9 @@ class Model(nn.Module): else: self._load(model, task=task) + # Delete super().training for accessing self.model.training + del self.training + def __call__( self, source: Union[str, Path, int, Image.Image, list, tuple, np.ndarray, torch.Tensor] = None, @@ -1143,3 +1146,29 @@ class Model(nn.Module): """ self.model.eval() return self + + def __getattr__(self, name): + """ + Enables accessing model attributes directly through the Model class. + + This method provides a way to access attributes of the underlying model directly through the Model class + instance. It first checks if the requested attribute is 'model', in which case it returns the model from + the module dictionary. Otherwise, it delegates the attribute lookup to the underlying model. + + Args: + name (str): The name of the attribute to retrieve. + + Returns: + (Any): The requested attribute value. + + Raises: + AttributeError: If the requested attribute does not exist in the model. + + Examples: + >>> model = YOLO("yolo11n.pt") + >>> print(model.stride) + >>> print(model.task) + """ + if name == "model": + return self._modules["model"] + return getattr(self.model, name)