allow for fine tuning on new static head

Author: lucidrains

README.md

@@ -135,7 +135,35 @@ model = load_pretrained_model('preview')
 
 ## Fine-tuning (wip)
 
-This repository will also allow for easy fine-tuning of Enformer. For starters, the following example shows a single step for finetuning on contextual data (cell type, transcription factor, etc)
+This repository will also allow for easy fine-tuning of Enformer.
+
+Fine-tuning on new tracks
+
+```python
+import torch
+from enformer_pytorch import Enformer
+from enformer_pytorch.finetune import HeadAdapterWrapper
+
+enformer = Enformer(
+    dim = 1536,
+    depth = 1,
+    heads = 8,
+    target_length = 200,
+)
+
+model = HeadAdapterWrapper(
+    enformer = enformer,
+    num_tracks = 128
+).cuda()
+
+seq = torch.randint(0, 4, (1, 196_608 // 2,)).cuda()
+target = torch.randn(1, 200, 128).cuda()  # 128 tracks
+
+loss = model(seq, target = target)
+loss.backward()
+```
+
+Finetuning on contextual data (cell type, transcription factor, etc)
 
 ```python
 import torch
@@ -151,7 +179,6 @@ enformer = Enformer(
 
 model = ContextAdapterWrapper(
     enformer = enformer,
-    enformer_dim = 1536,
     context_dim = 1024
 ).cuda()
 

enformer_pytorch/enformer_pytorch.py

@@ -265,6 +265,7 @@ def __init__(
         pos_dropout = 0.01
     ):
         super().__init__()
+        self.dim = dim
         self.num_alphabet = num_alphabet
         half_dim = dim // 2
         twice_dim = dim * 2

enformer_pytorch/finetune.py

@@ -1,5 +1,6 @@
 import torch
 from contextlib import contextmanager
+import torch.nn.functional as F
 from torch import nn, einsum
 from einops import rearrange
 from enformer_pytorch.enformer_pytorch import Enformer, poisson_loss
@@ -26,19 +27,56 @@ def freeze_batchnorm_context(model):
         for p, requires_grad in zip(bn.parameters(), state['requires_grad']):
             p.requires_grad = requires_grad
 
+class HeadAdapterWrapper(nn.Module):
+    def __init__(
+        self,
+        *,
+        enformer,
+        num_tracks
+    ):
+        super().__init__()
+        assert isinstance(enformer, Enformer)
+        self.enformer = enformer
+
+        self.to_tracks = nn.Sequential(
+            nn.Linear(enformer.dim * 2, num_tracks),
+            nn.Softplus()
+        )
+
+    def forward(
+        self,
+        seq,
+        *,
+        target = None,
+        freeze_enformer = False
+    ):
+        enformer_context = freeze_batchnorm_context(self.enformer) if not freeze_enformer else torch.no_grad()
+
+        with enformer_context:
+            _, embeddings = self.enformer(seq, return_embeddings = True)
+
+            if freeze_enformer:
+                embeddings.detach_()
+
+        preds = self.to_tracks(embeddings)
+
+        if not exists(target):
+            return preds
+
+        return poisson_loss(preds, target)
+
 class ContextAdapterWrapper(nn.Module):
     def __init__(
         self,
         *,
         enformer,
-        enformer_dim,
         context_dim
     ):
         super().__init__()
         assert isinstance(enformer, Enformer)
         self.enformer = enformer
 
-        self.to_context_weights = nn.Parameter(torch.randn(context_dim, enformer_dim * 2))
+        self.to_context_weights = nn.Parameter(torch.randn(context_dim, enformer.dim * 2))
         self.to_context_bias = nn.Parameter(torch.randn(context_dim))
 
     def forward(
@@ -62,6 +100,8 @@ def forward(
 
         pred = einsum('b n d, t d -> b n t', embeddings, weights) + bias
 
+        pred = F.softplus(pred)
+
         if not exists(target):
             return pred
 

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'enformer-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.1.4',
+  version = '0.1.5',
   license='MIT',
   description = 'Enformer - Pytorch',
   author = 'Phil Wang',