Diverge abstract class

1pha · Mar 5, 2024 · 0913fe2 · 0913fe2
1 parent a80b466
commit 0913fe2
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 4 deletions.
diff --git a/config/model/convnext_cls.yaml b/config/model/convnext_cls.yaml
@@ -1,4 +1,4 @@
-_target_: sage.models.base.ConvNext
+_target_: sage.models.base.ConvNextCls
 backbone:
   _target_: sage.models.model_zoo.convnext.build_convnext
   model_name: convnext-base

diff --git a/config/model/resnet_cls.yaml b/config/model/resnet_cls.yaml
@@ -1,4 +1,4 @@
-_target_: sage.models.base.ResNet
+_target_: sage.models.base.ResNetCls
 backbone:
   _target_: sage.models.model_zoo.resnet.build_resnet
   model_depth: 10

diff --git a/sage/models/base.py b/sage/models/base.py
@@ -53,15 +53,49 @@ def conv_layers(self):
             return find_conv_modules(self.backbone)
 
 
-class ResNet(ModelBase):
+class ClsBase(ModelBase):
+    def forward(self, brain: torch.Tensor, age: torch.Tensor):
+        pred = self.backbone(brain).squeeze()
+        loss = self.criterion(pred, age.long())
+        return dict(loss=loss,
+                    pred=pred.detach().cpu(),
+                    target=age.detach().cpu().long())
+
+
+class RegBase(ModelBase):
+    def forward(self, brain: torch.Tensor, age: torch.Tensor):
+        pred = self.backbone(brain).squeeze()
+        loss = self.criterion(pred, age.float())
+        return dict(loss=loss,
+                    pred=pred.detach().cpu(),
+                    target=age.detach().cpu())
+
+
+class ResNet(RegBase):
+    def __init__(self,
+                 backbone: nn.Module,
+                 criterion: nn.Module,
+                 name: str):
+        super().__init__(backbone=backbone, criterion=criterion, name=name)
+
+
+class ConvNext(RegBase):
+    def __init__(self,
+                 backbone: nn.Module,
+                 criterion: nn.Module,
+                 name: str):
+        super().__init__(backbone=backbone, criterion=criterion, name=name)
+
+
+class ResNetCls(RegBase):
     def __init__(self,
                  backbone: nn.Module,
                  criterion: nn.Module,
                  name: str):
         super().__init__(backbone=backbone, criterion=criterion, name=name)
 
 
-class ConvNext(ModelBase):
+class ConvNextCls(RegBase):
     def __init__(self,
                  backbone: nn.Module,
                  criterion: nn.Module,