Update models/deep_learning/backbone_model.py

models/deep_learning/backbone_model.py

@@ -1,109 +1,109 @@
-import os
-import sys
-
-sys.path.append(os.path.join(os.path.dirname(__file__), "..", ".."))
-
-import timm
-import torch
-import torch.nn as nn
-from transformers import CLIPModel as CLIPTransformersModel
-
-from utils import configs
-from utils.functional import check_data_type_variable, get_device
-
-
-class CLIPModel(nn.Module):
-    def __init__(
-        self,
-        model_clip_name: str,
-        freeze_model: bool,
-        pretrained_model: bool,
-        num_classes: int,
-    ):
-        super().__init__()
-        self.model_clip_name = model_clip_name
-        self.freeze_model = freeze_model
-        self.pretrained_model = pretrained_model
-        self.num_classes = num_classes
-        self.device = get_device()
-
-        self.check_arguments()
-        self.init_model()
-
-    def check_arguments(self):
-        check_data_type_variable(self.model_clip_name, str)
-        check_data_type_variable(self.freeze_model, bool)
-        check_data_type_variable(self.pretrained_model, bool)
-        check_data_type_variable(self.num_classes, int)
-
-        if self.model_clip_name != configs.CLIP_NAME_MODEL:
-            raise ValueError(
-                f"Model clip name must be {configs.CLIP_NAME_MODEL}, but it is {self.model_clip_name}"
-            )
-
-    def init_model(self):
-        clip_model = CLIPTransformersModel.from_pretrained(self.model_clip_name)
-        for layer in clip_model.children():
-            if hasattr(layer, "reset_parameters") and not self.pretrained_model:
-                layer.reset_parameters()
-        for param in clip_model.parameters():
-            param.required_grad = False if not self.freeze_model else True
-        self.vision_model = clip_model.vision_model.to(self.device)
-        self.visual_projection = clip_model.visual_projection.to(self.device).to(
-            self.device
-        )
-        self.classifier = nn.Linear(
-            512, 1 if self.num_classes in (1, 2) else self.num_classes
-        ).to(self.device)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        x = self.vision_model(x)
-        x = self.visual_projection(x.pooler_output)
-        x = self.classifier(x)
-        return x
-
-
-class TorchModel(nn.Module):
-    def __init__(
-        self,
-        name_model: str,
-        freeze_model: bool,
-        pretrained_model: bool,
-        num_classes: int,
-    ):
-        super().__init__()
-        self.name_model = name_model
-        self.freeze_model = freeze_model
-        self.pretrained_model = pretrained_model
-        self.num_classes = num_classes
-        self.device = get_device()
-
-        self.check_arguments()
-        self.init_model()
-
-    def check_arguments(self):
-        check_data_type_variable(self.name_model, str)
-        check_data_type_variable(self.freeze_model, bool)
-        check_data_type_variable(self.pretrained_model, bool)
-        check_data_type_variable(self.num_classes, int)
-
-        if self.name_model not in tuple(configs.NAME_MODELS.keys()):
-            raise ValueError(
-                f"Name model must be in {tuple(configs.NAME_MODELS.keys())}, but it is {self.name_model}"
-            )
-
-    def init_model(self):
-        self.model = timm.create_model(
-            self.name_model, pretrained=self.pretrained_model, num_classes=0
-        ).to(self.device)
-        for param in self.model.parameters():
-            param.required_grad = False if not self.freeze_model else True
-        self.classifier = nn.Linear(
-            self.model.num_features,
-            1 if self.num_classes in (1, 2) else self.num_classes,
-        ).to(self.device)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        x = self.model(x)
-        x = self.classifier(x)
-        return x
+import os
+import sys
+
+sys.path.append(os.path.join(os.path.dirname(__file__), "..", ".."))
+
+import timm
+import torch
+import torch.nn as nn
+from transformers import CLIPModel as CLIPTransformersModel
+
+from utils import configs
+from utils.functional import check_data_type_variable, get_device
+
+
+class CLIPModel(nn.Module):
+    def __init__(
+        self,
+        model_clip_name: str,
+        freeze_model: bool,
+        pretrained_model: bool,
+        num_classes: int,
+    ):
+        super().__init__()
+        self.model_clip_name = model_clip_name
+        self.freeze_model = freeze_model
+        self.pretrained_model = pretrained_model
+        self.num_classes = num_classes
+        self.device = get_device()
+
+        self.check_arguments()
+        self.init_model()
+
+    def check_arguments(self):
+        check_data_type_variable(self.model_clip_name, str)
+        check_data_type_variable(self.freeze_model, bool)
+        check_data_type_variable(self.pretrained_model, bool)
+        check_data_type_variable(self.num_classes, int)
+
+        if self.model_clip_name != configs.CLIP_NAME_MODEL:
+            raise ValueError(
+                f"Model clip name must be {configs.CLIP_NAME_MODEL}, but it is {self.model_clip_name}"
+            )
+
+    def init_model(self):
+        self.clip_model = CLIPTransformersModel.from_pretrained(self.model_clip_name)
+        for layer in self.clip_model.children():
+            if hasattr(layer, "reset_parameters") and not self.pretrained_model:
+                layer.reset_parameters()
+        for param in self.clip_model.parameters():
+            param.required_grad = False if not self.freeze_model else True
+        self.vision_model = self.clip_model.vision_model.to(self.device)
+        self.visual_projection = self.clip_model.visual_projection.to(self.device).to(
+            self.device
+        )
+        self.classifier = nn.Linear(
+            512, 1 if self.num_classes in (1, 2) else self.num_classes
+        ).to(self.device)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.vision_model(x)
+        x = self.visual_projection(x.pooler_output)
+        x = self.classifier(x)
+        return x
+
+
+class TorchModel(nn.Module):
+    def __init__(
+        self,
+        name_model: str,
+        freeze_model: bool,
+        pretrained_model: bool,
+        num_classes: int,
+    ):
+        super().__init__()
+        self.name_model = name_model
+        self.freeze_model = freeze_model
+        self.pretrained_model = pretrained_model
+        self.num_classes = num_classes
+        self.device = get_device()
+
+        self.check_arguments()
+        self.init_model()
+
+    def check_arguments(self):
+        check_data_type_variable(self.name_model, str)
+        check_data_type_variable(self.freeze_model, bool)
+        check_data_type_variable(self.pretrained_model, bool)
+        check_data_type_variable(self.num_classes, int)
+
+        if self.name_model not in tuple(configs.NAME_MODELS.keys()):
+            raise ValueError(
+                f"Name model must be in {tuple(configs.NAME_MODELS.keys())}, but it is {self.name_model}"
+            )
+
+    def init_model(self):
+        self.model = timm.create_model(
+            self.name_model, pretrained=self.pretrained_model, num_classes=0
+        ).to(self.device)
+        for param in self.model.parameters():
+            param.required_grad = False if not self.freeze_model else True
+        self.classifier = nn.Linear(
+            self.model.num_features,
+            1 if self.num_classes in (1, 2) else self.num_classes,
+        ).to(self.device)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.model(x)
+        x = self.classifier(x)
+        return x