feat: ✨ Add activation specification for UNets/ScaleNets

janelia-cellmap · May 2, 2024 · 1f22595 · 1f22595
1 parent 1449ca3
commit 1f22595
Show file tree

Hide file tree

Showing 4 changed files with 17 additions and 2 deletions.
diff --git a/src/leibnetz/leibnet.py b/src/leibnetz/leibnet.py
@@ -19,7 +19,7 @@ def __init__(
         self,
         nodes: Iterable,
         outputs: dict[str, Sequence[Tuple]],
-        retain_buffer=True,
+        # retain_buffer=True,
         initialization="kaiming",
         name="LeibNet",
     ):
@@ -73,7 +73,7 @@ def __init__(
             pass
         else:
             raise ValueError(f"Unknown initialization {initialization}")
-        self.retain_buffer = retain_buffer
+        # self.retain_buffer = retain_buffer
         self.retain_buffer = True
         if torch.cuda.is_available():
             self.cuda()

diff --git a/src/leibnetz/nets/attentive_scalenet.py b/src/leibnetz/nets/attentive_scalenet.py
@@ -22,6 +22,7 @@ def build_subnet(
     norm_layer=None,
     residual=False,
     dropout_prob=None,
+    activation="ReLU",
 ):
     # define downsample nodes
     downsample_factors = np.array(downsample_factors)
@@ -42,6 +43,7 @@ def build_subnet(
                 norm_layer=norm_layer,
                 residual=residual,
                 dropout_prob=dropout_prob,
+                activation=activation,
             ),
         )
         c += 1
@@ -76,6 +78,7 @@ def build_subnet(
             norm_layer=norm_layer,
             residual=residual,
             dropout_prob=dropout_prob,
+            activation=activation,
         )
     )
     input_key = output_key
@@ -121,6 +124,7 @@ def build_subnet(
                 norm_layer=norm_layer,
                 residual=residual,
                 dropout_prob=dropout_prob,
+                activation=activation,
             )
         )
         input_key = output_key
@@ -138,6 +142,7 @@ def build_subnet(
             norm_layer=norm_layer,
             residual=residual,
             dropout_prob=dropout_prob,
+            activation=activation,
         )
     )
     outputs = {

diff --git a/src/leibnetz/nets/scalenet.py b/src/leibnetz/nets/scalenet.py
@@ -18,6 +18,7 @@ def build_subnet(
     norm_layer=None,
     residual=False,
     dropout_prob=None,
+    activation="ReLU",
 ):
     # define downsample nodes
     downsample_factors = np.array(downsample_factors)
@@ -37,6 +38,7 @@ def build_subnet(
                 norm_layer=norm_layer,
                 residual=residual,
                 dropout_prob=dropout_prob,
+                activation=activation,
             ),
         )
         c += 1
@@ -71,6 +73,7 @@ def build_subnet(
             norm_layer=norm_layer,
             residual=residual,
             dropout_prob=dropout_prob,
+            activation=activation,
         )
     )
     input_key = output_key
@@ -101,6 +104,7 @@ def build_subnet(
                 norm_layer=norm_layer,
                 residual=residual,
                 dropout_prob=dropout_prob,
+                activation=activation,
             )
         )
         input_key = output_key
@@ -118,6 +122,7 @@ def build_subnet(
             norm_layer=norm_layer,
             residual=residual,
             dropout_prob=dropout_prob,
+            activation=activation,
         )
     )
     outputs = {

diff --git a/src/leibnetz/nets/unet.py b/src/leibnetz/nets/unet.py
@@ -15,6 +15,7 @@ def build_unet(
     norm_layer=None,
     residual=False,
     dropout_prob=None,
+    activation="ReLU",
 ):
     # define downsample nodes
     downsample_factors = np.array(downsample_factors)
@@ -34,6 +35,7 @@ def build_unet(
                 norm_layer=norm_layer,
                 residual=residual,
                 dropout_prob=dropout_prob,
+                activation=activation,
             ),
         )
         c += 1
@@ -62,6 +64,7 @@ def build_unet(
             norm_layer=norm_layer,
             residual=residual,
             dropout_prob=dropout_prob,
+            activation=activation,
         )
     )
     input_key = output_key
@@ -92,6 +95,7 @@ def build_unet(
                 norm_layer=norm_layer,
                 residual=residual,
                 dropout_prob=dropout_prob,
+                activation=activation,
             )
         )
         input_key = output_key
@@ -109,6 +113,7 @@ def build_unet(
             norm_layer=norm_layer,  # TODO: remove?
             residual=residual,
             dropout_prob=dropout_prob,
+            activation=activation,
         )
     )