Update model.py

model.py

@@ -1,144 +1,142 @@
-import torch
-import torch.nn as nn
-import torchvision.models as models
-from modelscope.msdatasets import MsDataset
-from utils import MODEL_DIR
-
-
-def get_backbone(ver, backbone_list):
-    for bb in backbone_list:
-        if ver == bb["ver"]:
-            return bb
-
-    print("Backbone name not found, using default option - alexnet.")
-    return backbone_list[0]
-
-
-def model_info(m_ver):
-    backbone_list = MsDataset.load(
-        "monetjoe/cv_backbones", subset_name="ImageNet1k_v1", split="train"
-    )
-    backbone = get_backbone(m_ver, backbone_list)
-    m_type = str(backbone["type"])
-    input_size = int(backbone["input_size"])
-    return m_type, input_size
-
-
-def Classifier(cls_num: int, output_size: int, linear_output: bool):
-    q = (1.0 * output_size / cls_num) ** 0.25
-    l1 = int(q * cls_num)
-    l2 = int(q * l1)
-    l3 = int(q * l2)
-
-    if linear_output:
-        return torch.nn.Sequential(
-            nn.Dropout(),
-            nn.Linear(output_size, l3),
-            nn.ReLU(inplace=True),
-            nn.Dropout(),
-            nn.Linear(l3, l2),
-            nn.ReLU(inplace=True),
-            nn.Dropout(),
-            nn.Linear(l2, l1),
-            nn.ReLU(inplace=True),
-            nn.Linear(l1, cls_num),
-        )
-
-    else:
-        return torch.nn.Sequential(
-            nn.Dropout(),
-            nn.Conv2d(output_size, l3, kernel_size=(1, 1), stride=(1, 1)),
-            nn.ReLU(inplace=True),
-            nn.AdaptiveAvgPool2d(output_size=(1, 1)),
-            nn.Flatten(),
-            nn.Linear(l3, l2),
-            nn.ReLU(inplace=True),
-            nn.Dropout(),
-            nn.Linear(l2, l1),
-            nn.ReLU(inplace=True),
-            nn.Linear(l1, cls_num),
-        )
-
-
-class EvalNet:
-    model = None
-    m_type = "squeezenet"
-    input_size = 224
-    output_size = 512
-
-    def __init__(self, log_name: str, cls_num=4):
-        saved_model_path = f"{MODEL_DIR}/{log_name}/save.pt"
-        m_ver = "_".join(log_name.split("_")[:-1])
-        self.m_type, self.input_size = model_info(m_ver)
-
-        if not hasattr(models, m_ver):
-            print("Unsupported model.")
-            exit()
-
-        self.model = eval("models.%s()" % m_ver)
-        linear_output = self._set_outsize()
-        self._set_classifier(cls_num, linear_output)
-        checkpoint = torch.load(saved_model_path, map_location="cpu")
-        if torch.cuda.is_available():
-            checkpoint = torch.load(saved_model_path)
-
-        self.model.load_state_dict(checkpoint, False)
-        self.model.eval()
-
-    def _set_outsize(self, debug_mode=False):
-        for name, module in self.model.named_modules():
-            if (
-                str(name).__contains__("classifier")
-                or str(name).__eq__("fc")
-                or str(name).__contains__("head")
-            ):
-                if isinstance(module, torch.nn.Linear):
-                    self.output_size = module.in_features
-                    if debug_mode:
-                        print(
-                            f"{name}(Linear): {self.output_size} -> {module.out_features}"
-                        )
-                    return True
-
-                if isinstance(module, torch.nn.Conv2d):
-                    self.output_size = module.in_channels
-                    if debug_mode:
-                        print(
-                            f"{name}(Conv2d): {self.output_size} -> {module.out_channels}"
-                        )
-                    return False
-
-        return False
-
-    def _set_classifier(self, cls_num, linear_output):
-        if self.m_type == "convnext":
-            del self.model.classifier[2]
-            self.model.classifier = nn.Sequential(
-                *list(self.model.classifier)
-                + list(Classifier(cls_num, self.output_size, linear_output))
-            )
-            return
-
-        if hasattr(self.model, "classifier"):
-            self.model.classifier = Classifier(cls_num, self.output_size, linear_output)
-            return
-
-        elif hasattr(self.model, "fc"):
-            self.model.fc = Classifier(cls_num, self.output_size, linear_output)
-            return
-
-        elif hasattr(self.model, "head"):
-            self.model.head = Classifier(cls_num, self.output_size, linear_output)
-            return
-
-        self.model.heads.head = Classifier(cls_num, self.output_size, linear_output)
-
-    def forward(self, x):
-        if torch.cuda.is_available():
-            x = x.cuda()
-            self.model = self.model.cuda()
-
-        if self.m_type == "googlenet" and self.training:
-            return self.model(x)[0]
-        else:
-            return self.model(x)
+import torch
+import torch.nn as nn
+import torchvision.models as models
+from modelscope.msdatasets import MsDataset
+from utils import MODEL_DIR
+
+
+def get_backbone(ver, backbone_list):
+    for bb in backbone_list:
+        if ver == bb["ver"]:
+            return bb
+
+    print("Backbone name not found, using default option - alexnet.")
+    return backbone_list[0]
+
+
+def model_info(m_ver):
+    backbone_list = MsDataset.load("monetjoe/cv_backbones", split="train")
+    backbone = get_backbone(m_ver, backbone_list)
+    m_type = str(backbone["type"])
+    input_size = int(backbone["input_size"])
+    return m_type, input_size
+
+
+def Classifier(cls_num: int, output_size: int, linear_output: bool):
+    q = (1.0 * output_size / cls_num) ** 0.25
+    l1 = int(q * cls_num)
+    l2 = int(q * l1)
+    l3 = int(q * l2)
+
+    if linear_output:
+        return torch.nn.Sequential(
+            nn.Dropout(),
+            nn.Linear(output_size, l3),
+            nn.ReLU(inplace=True),
+            nn.Dropout(),
+            nn.Linear(l3, l2),
+            nn.ReLU(inplace=True),
+            nn.Dropout(),
+            nn.Linear(l2, l1),
+            nn.ReLU(inplace=True),
+            nn.Linear(l1, cls_num),
+        )
+
+    else:
+        return torch.nn.Sequential(
+            nn.Dropout(),
+            nn.Conv2d(output_size, l3, kernel_size=(1, 1), stride=(1, 1)),
+            nn.ReLU(inplace=True),
+            nn.AdaptiveAvgPool2d(output_size=(1, 1)),
+            nn.Flatten(),
+            nn.Linear(l3, l2),
+            nn.ReLU(inplace=True),
+            nn.Dropout(),
+            nn.Linear(l2, l1),
+            nn.ReLU(inplace=True),
+            nn.Linear(l1, cls_num),
+        )
+
+
+class EvalNet:
+    model = None
+    m_type = "squeezenet"
+    input_size = 224
+    output_size = 512
+
+    def __init__(self, log_name: str, cls_num=4):
+        saved_model_path = f"{MODEL_DIR}/{log_name}/save.pt"
+        m_ver = "_".join(log_name.split("_")[:-1])
+        self.m_type, self.input_size = model_info(m_ver)
+
+        if not hasattr(models, m_ver):
+            print("Unsupported model.")
+            exit()
+
+        self.model = eval("models.%s()" % m_ver)
+        linear_output = self._set_outsize()
+        self._set_classifier(cls_num, linear_output)
+        checkpoint = torch.load(saved_model_path, map_location="cpu")
+        if torch.cuda.is_available():
+            checkpoint = torch.load(saved_model_path)
+
+        self.model.load_state_dict(checkpoint, False)
+        self.model.eval()
+
+    def _set_outsize(self, debug_mode=False):
+        for name, module in self.model.named_modules():
+            if (
+                str(name).__contains__("classifier")
+                or str(name).__eq__("fc")
+                or str(name).__contains__("head")
+            ):
+                if isinstance(module, torch.nn.Linear):
+                    self.output_size = module.in_features
+                    if debug_mode:
+                        print(
+                            f"{name}(Linear): {self.output_size} -> {module.out_features}"
+                        )
+                    return True
+
+                if isinstance(module, torch.nn.Conv2d):
+                    self.output_size = module.in_channels
+                    if debug_mode:
+                        print(
+                            f"{name}(Conv2d): {self.output_size} -> {module.out_channels}"
+                        )
+                    return False
+
+        return False
+
+    def _set_classifier(self, cls_num, linear_output):
+        if self.m_type == "convnext":
+            del self.model.classifier[2]
+            self.model.classifier = nn.Sequential(
+                *list(self.model.classifier)
+                + list(Classifier(cls_num, self.output_size, linear_output))
+            )
+            return
+
+        if hasattr(self.model, "classifier"):
+            self.model.classifier = Classifier(cls_num, self.output_size, linear_output)
+            return
+
+        elif hasattr(self.model, "fc"):
+            self.model.fc = Classifier(cls_num, self.output_size, linear_output)
+            return
+
+        elif hasattr(self.model, "head"):
+            self.model.head = Classifier(cls_num, self.output_size, linear_output)
+            return
+
+        self.model.heads.head = Classifier(cls_num, self.output_size, linear_output)
+
+    def forward(self, x):
+        if torch.cuda.is_available():
+            x = x.cuda()
+            self.model = self.model.cuda()
+
+        if self.m_type == "googlenet" and self.training:
+            return self.model(x)[0]
+        else:
+            return self.model(x)