add ability to use the embeddings from the output of the transformer block, and not from after the final pointwise, use layernorm if taking embeddings from post-transformer

Author: lucidrains

README.md

@@ -180,7 +180,8 @@ enformer = Enformer.from_hparams(
     
 model = HeadAdapterWrapper(
     enformer = enformer,
-    num_tracks = 128
+    num_tracks = 128,
+    post_transformer_embed = False   # by default, embeddings are taken from after the final pointwise block w/ conv -> gelu - but if you'd like the embeddings right after the transformer block with a learned layernorm, set this to True
 ).cuda()
 
 seq = torch.randint(0, 5, (1, 196_608 // 2,)).cuda()

enformer_pytorch/finetune.py

@@ -1,4 +1,5 @@
 import torch
+from copy import deepcopy
 from contextlib import contextmanager
 import torch.nn.functional as F
 from torch import nn, einsum
@@ -16,6 +17,11 @@ def exists(val):
 def null_context():
     yield
 
+# better sequential
+
+def Sequential(*modules):
+    return nn.Sequential(*filter(exists, modules))
+
 # controlling freezing of layers
 
 def set_module_requires_grad_(module, requires_grad):
@@ -88,14 +94,16 @@ def __init__(
         *,
         enformer,
         num_tracks,
+        post_transformer_embed = False, # whether to take the embeddings from right after the transformer, instead of after the final pointwise convolutional - this would add another layernorm
         discrete_key_value_bottleneck = False,
         bottleneck_num_memories = 256,
         bottleneck_num_codebooks = 4,
         bottleneck_decay = 0.9,
+        transformer_embed_fn: nn.Module = nn.Identity()
     ):
         super().__init__()
         assert isinstance(enformer, Enformer)
-        enformer_hidden_dim = enformer.dim * 2
+        enformer_hidden_dim = enformer.dim * (2 if not post_transformer_embed else 1)
 
         self.discrete_key_value_bottleneck = discrete_key_value_bottleneck
 
@@ -109,8 +117,20 @@ def __init__(
                 decay = bottleneck_decay,
             )
 
+        self.post_transformer_embed = post_transformer_embed
+
         self.enformer = enformer
 
+        if post_transformer_embed:
+            self.enformer = deepcopy(enformer)
+            self.enformer._trunk[-1] = nn.Identity()
+            self.enformer.final_pointwise = nn.Identity()
+
+        self.post_embed_transform = Sequential(
+            transformer_embed_fn,
+            nn.LayerNorm(enformer_hidden_dim) if post_transformer_embed else None
+        )
+
         self.to_tracks = nn.Sequential(
             nn.Linear(enformer_hidden_dim, num_tracks),
             nn.Softplus()

setup.py

@@ -4,7 +4,7 @@
   name = 'enformer-pytorch',
   packages = find_packages(exclude=[]),
   include_package_data = True,
-  version = '0.5.7',
+  version = '0.6.0',
   license='MIT',
   description = 'Enformer - Pytorch',
   author = 'Phil Wang',