Upload directory

models/base/__init__.py

@@ -0,0 +1,151 @@
+import os
+from typing import Union
+import torch
+from torch import device
+from .utils import get_parameter_device, get_parameter_dtype, save_state_dict_and_config, load_state_dict_from_path
+
+class BaseModel(torch.nn.Module):
+    """
+    A base model class that provides a template for implementing models. It includes methods for
+    loading, saving, and managing model configurations and states. This class is designed to be
+    extended by specific model implementations.
+
+    Attributes:
+        config (object): Configuration object containing model settings.
+        input_color_flip (bool): Whether to flip the color channels from BGR to RGB.
+    """
+
+    def __init__(self, config=None):
+        """
+        Initializes the BaseModel class.
+
+        Parameters:
+            config (object, optional): Configuration object containing model settings.
+        """
+        super(BaseModel, self).__init__()
+        self.config = config
+        if self.config.color_space == 'BGR':
+            self.input_color_flip = True
+            self._config_color_space = 'BGR'
+            self.config.color_space = 'RGB'
+        else:
+            self.input_color_flip = False
+
+    def forward(self, x):
+        """
+        Forward pass of the model. Needs to be implemented in subclass.
+
+        Parameters:
+            x (torch.Tensor): Input tensor.
+
+        Raises:
+            NotImplementedError: If the subclass does not implement this method.
+        """
+        raise NotImplementedError('forward must be implemented in subclass')
+
+    @classmethod
+    def from_config(cls, config) -> "BaseModel":
+        """
+        Creates an instance of this class from a configuration object. Needs to be implemented in subclass.
+
+        Parameters:
+            config (object): Configuration object.
+
+        Returns:
+            BaseModel: An instance of the subclass.
+
+        Raises:
+            NotImplementedError: If the subclass does not implement this method.
+        """
+        raise NotImplementedError('from_config must be implemented in subclass')
+
+    def make_train_transform(self):
+        """
+        Creates training data transformations. Needs to be implemented in subclass.
+
+        Raises:
+            NotImplementedError: If the subclass does not implement this method.
+        """
+        raise NotImplementedError('make_train_transform must be implemented in subclass')
+
+    def make_test_transform(self):
+        """
+        Creates testing data transformations. Needs to be implemented in subclass.
+
+        Raises:
+            NotImplementedError: If the subclass does not implement this method.
+        """
+        raise NotImplementedError('make_test_transform must be implemented in subclass')
+
+    def save_pretrained(
+        self,
+        save_dir: Union[str, os.PathLike],
+        name: str = 'model.pt',
+        rank: int = 0,
+    ):
+        """
+        Saves the model's state_dict and configuration to the specified directory.
+
+        Parameters:
+            save_dir (Union[str, os.PathLike]): The directory to save the model.
+            name (str, optional): The name of the file to save the model as. Default is 'model.pt'.
+            rank (int, optional): The rank of the process (used in distributed training). Default is 0.
+        """
+        save_path = os.path.join(save_dir, name)
+        if rank == 0:
+            save_state_dict_and_config(self.state_dict(), self.config, save_path)
+
+    def load_state_dict_from_path(self, pretrained_model_path):
+        state_dict = load_state_dict_from_path(pretrained_model_path)
+        if 'net.vit' in list(self.state_dict().keys())[-1] and 'pretrained_models' in pretrained_model_path:
+            state_dict = {k.replace('net', 'net.vit'): v for k, v in state_dict.items()}
+
+        st_keys = list(state_dict.keys())
+        self_keys = list(self.state_dict().keys())
+        print('compatible keys in state_dict', len(set(st_keys).intersection(set(self_keys))), '/', len(st_keys))
+        print('Check\n\n')
+        result = self.load_state_dict(state_dict, strict=False)
+        print(result)
+        print(f"Loaded pretrained model from {pretrained_model_path}")
+
+
+    @property
+    def device(self) -> device:
+        """
+        Returns the device of the model's parameters.
+
+        Returns:
+            device: The device the model is on.
+        """
+        return get_parameter_device(self)
+
+    @property
+    def dtype(self) -> torch.dtype:
+        """
+        Returns the data type of the model's parameters.
+
+        Returns:
+            torch.dtype: The data type of the model.
+        """
+        return get_parameter_dtype(self)
+
+    def num_parameters(self, only_trainable: bool = False) -> int:
+        """
+        Returns the number of parameters in the model, optionally filtering only trainable parameters.
+
+        Parameters:
+            only_trainable (bool, optional): Whether to count only trainable parameters. Default is False.
+
+        Returns:
+            int: The number of parameters.
+        """
+        return sum(p.numel() for p in self.parameters() if p.requires_grad or not only_trainable)
+
+    def has_trainable_params(self):
+        """
+        Checks if the model has any trainable parameters.
+
+        Returns:
+            bool: True if the model has trainable parameters, False otherwise.
+        """
+        return any(p.requires_grad for p in self.parameters())