add ConvCore for convolutional models only

neuralpredictors/layers/cores/base.py

@@ -9,10 +9,6 @@ class Core(ABC):
     Base class for the core models, taking 2d inputs and computing nonlinear features.
     """
 
-    def __init__(self) -> None:
-        super().__init__()
-        self.set_batchnorm_type()
-
     def initialize(self):
         """
         Initialization applied on the core.
@@ -34,6 +30,35 @@ def init_conv(m):
             if m.bias is not None:
                 m.bias.data.fill_(0)
 
+    @abstractmethod
+    def regularizer(self):
+        """
+        Regularization applied on the core. Returns a scalar value.
+        """
+
+    @abstractmethod
+    def forward(self, x):
+        """
+        Forward function for pytorch nn module.
+
+        Args:
+            x (torch.tensor): input of shape (batch, channels, height, width)
+        """
+
+    def __repr__(self):
+        s = super().__repr__()
+        s += f" [{self.__class__.__name__} regularizers: "
+        ret = []
+        for attr in filter(lambda x: "gamma" in x or "skip" in x, dir(self)):
+            ret.append(f"{attr} = {getattr(self, attr)}")
+        return s + "|".join(ret) + "]\n"
+
+
+class ConvCore(Core):
+    def __init__(self) -> None:
+        super().__init__()
+        self.set_batchnorm_type()
+
     @abstractmethod
     def set_batchnorm_type(self):
         """
@@ -63,26 +88,3 @@ def add_bn_layer(self, layer: OrderedDict, layer_idx: int):
                 layer["bias"] = self.bias_layer_cls(hidden_channels)
             elif not bias and scale:
                 layer["scale"] = self.scale_layer_cls(hidden_channels)
-
-    @abstractmethod
-    def regularizer(self):
-        """
-        Regularization applied on the core. Returns a scalar value.
-        """
-
-    @abstractmethod
-    def forward(self, x):
-        """
-        Forward function for pytorch nn module.
-
-        Args:
-            x (torch.tensor): input of shape (batch, channels, height, width)
-        """
-
-    def __repr__(self):
-        s = super().__repr__()
-        s += f" [{self.__class__.__name__} regularizers: "
-        ret = []
-        for attr in filter(lambda x: "gamma" in x or "skip" in x, dir(self)):
-            ret.append(f"{attr} = {getattr(self, attr)}")
-        return s + "|".join(ret) + "]\n"

neuralpredictors/layers/cores/conv2d.py

@@ -27,12 +27,12 @@
     RotationEquivariantScale2DLayer,
 )
 from ..squeeze_excitation import SqueezeExcitationBlock
-from .base import Core
+from .base import ConvCore, Core
 
 logger = logging.getLogger(__name__)
 
 
-class Stacked2dCore(Core, nn.Module):
+class Stacked2dCore(ConvCore, nn.Module):
     """
     An instantiation of the Core base class. Made up of layers layers of nn.sequential modules.
     Allows for the flexible implementations of many different architectures, such as convolutional layers,

neuralpredictors/layers/cores/conv3d.py

@@ -11,10 +11,10 @@
 
 from ...regularizers import DepthLaplaceL21d
 from ..affine import Bias3DLayer, Scale3DLayer
-from .base import Core
+from .base import ConvCore
 
 
-class Core3d(Core):
+class Core3d(ConvCore):
     def initialize(self, cuda=False):
         self.apply(self.init_conv)
         self.put_to_cuda(cuda=cuda)