add transfer modules

src/baskerville/blocks.py

@@ -195,7 +195,7 @@ def conv_dna(
         kernel_initializer=kernel_initializer,
         kernel_regularizer=tf.keras.regularizers.l2(l2_scale),
     )(current)
-
+   
     # squeeze-excite
     if se:
         current = squeeze_excite(current)
@@ -1109,6 +1109,9 @@ def transformer(
     qkv_width=1,
     mha_initializer="he_normal",
     kernel_initializer="he_normal",
+    adapter=None,
+    latent=16,
+    seqlen_train=None,
     **kwargs,
 ):
     """Construct a transformer block.
@@ -1140,20 +1143,25 @@ def transformer(
         initializer=mha_initializer,
         l2_scale=mha_l2_scale,
         qkv_width=qkv_width,
+        seqlen_train=seqlen_train
     )(current)
 
     # dropout
     if dropout > 0:
         current = tf.keras.layers.Dropout(dropout)(current)
 
+    # add houlsby-adapter
+    if adapter=='houlsby':
+        current = layers.AdapterHoulsby(latent_size=latent)(current)
+
     # residual
     current = tf.keras.layers.Add()([inputs, current])
 
     if dense_expansion == 0:
         final = current
     else:
         final = transformer_dense(
-            current, out_size, dense_expansion, l2_scale, dropout, kernel_initializer
+            current, out_size, dense_expansion, l2_scale, dropout, kernel_initializer, adapter, latent
         )
 
     return final
@@ -1265,7 +1273,8 @@ def transformer_split(
 
 
 def transformer_dense(
-    inputs, out_size, dense_expansion, l2_scale, dropout, kernel_initializer
+    inputs, out_size, dense_expansion, l2_scale, dropout, kernel_initializer,
+    adapter=None, latent=16
 ):
     """Transformer block dense portion."""
     # layer norm
@@ -1297,6 +1306,9 @@ def transformer_dense(
     if dropout > 0:
         current = tf.keras.layers.Dropout(dropout)(current)
 
+    if adapter=='houlsby':
+        current = layers.AdapterHoulsby(latent_size=latent)(current)
+
     # residual
     final = tf.keras.layers.Add()([inputs, current])
 
@@ -1439,11 +1451,20 @@ def squeeze_excite(
     additive=False,
     norm_type=None,
     bn_momentum=0.9,
+    kernel_initializer='glorot_uniform',
+    use_bias=True,
+    scale_fun='sigmoid',
     **kwargs,
 ):
     return layers.SqueezeExcite(
-        activation, additive, bottleneck_ratio, norm_type, bn_momentum
-    )(inputs)
+        activation=activation, 
+        additive=additive, 
+        bottleneck_ratio=bottleneck_ratio, 
+        norm_type=norm_type, 
+        bn_momentum=bn_momentum, 
+        kernel_initializer=kernel_initializer, 
+        scale_fun=scale_fun,
+        use_bias=use_bias)(inputs)
 
 
 def wheeze_excite(inputs, pool_size, **kwargs):