add se_adapter and locon

calico · Jul 3, 2024 · c5b6724 · c5b6724
1 parent cdf3634
commit c5b6724
Show file tree

Hide file tree

Showing 8 changed files with 538 additions and 533 deletions.
diff --git a/src/baskerville/HY_helper.py b/src/baskerville/HY_helper.py
diff --git a/src/baskerville/blocks.py b/src/baskerville/blocks.py
@@ -149,8 +149,6 @@ def conv_dna(
     conv_type="standard",
     kernel_initializer="he_normal",
     padding="same",
-    transfer_se=False,
-    se_ratio=16,
 ):
     """Construct a single convolution block, assumed to be operating on DNA.
 
@@ -197,19 +195,7 @@ def conv_dna(
         kernel_initializer=kernel_initializer,
         kernel_regularizer=tf.keras.regularizers.l2(l2_scale),
     )(current)
-
-    # squeeze-excite for transfer
-    if transfer_se:
-        se_out = squeeze_excite(current, 
-                                activation=None,
-                                additive=False, 
-                                bottleneck_ratio=se_ratio,
-                                use_bias=False,
-                                kernel_initializer=tf.keras.initializers.TruncatedNormal(stddev=1e-3),
-                                scale_fun='tanh'
-                               )
-        current = current + se_out
-
+
     # squeeze-excite
     if se:
         current = squeeze_excite(current)
@@ -281,8 +267,6 @@ def conv_nac(
     kernel_initializer="he_normal",
     padding="same",
     se=False,
-    transfer_se=False,
-    se_ratio=16,
 ):
     """Construct a single convolution block.
 
@@ -342,18 +326,6 @@ def conv_nac(
         kernel_regularizer=tf.keras.regularizers.l2(l2_scale),
     )(current)
 
-    # squeeze-excite for transfer
-    if transfer_se:
-        se_out = squeeze_excite(current, 
-                                activation=None,
-                                additive=False, 
-                                bottleneck_ratio=se_ratio,
-                                use_bias=False,
-                                kernel_initializer=tf.keras.initializers.TruncatedNormal(stddev=1e-3),
-                                scale_fun='tanh'                                
-                               )
-        current = current + se_out
-
     # squeeze-excite
     if se:
         current = squeeze_excite(current)
@@ -484,8 +456,6 @@ def unet_conv(
     bn_momentum=0.99,
     kernel_size=1,
     kernel_initializer="he_normal",
-    transfer_se=False,
-    se_ratio=16,
     upsample_conv=False,
 ):
     """Construct a feature pyramid network block.
@@ -561,17 +531,6 @@ def unet_conv(
         kernel_initializer=kernel_initializer,
     )(current)
 
-    if transfer_se:
-        se_out = squeeze_excite(current, 
-                                activation=None,
-                                additive=False, 
-                                bottleneck_ratio=se_ratio,
-                                use_bias=False,
-                                kernel_initializer=tf.keras.initializers.TruncatedNormal(stddev=1e-3),
-                                scale_fun='tanh'                                
-                               )
-        current = current + se_out
-
     # dropout
     if dropout > 0:
         current = tf.keras.layers.Dropout(dropout)(current)