add more freezing unfreezing functions for finetuning

Author: lucidrains

enformer_pytorch/finetune.py

@@ -14,10 +14,18 @@ def exists(val):
 def null_context():
     yield
 
+# controlling freezing of layers
+
 def set_module_requires_grad_(module, requires_grad):
     for param in module.parameters():
         param.requires_grad = requires_grad
 
+def freeze_all_layers_(module):
+    set_module_requires_grad_(module, False)
+
+def unfreeze_all_layers_(module):
+    set_module_requires_grad_(module, True)
+
 def freeze_batchnorms_(model):
     bns = [m for m in model.modules() if isinstance(m, nn.BatchNorm1d)]
 
@@ -30,6 +38,24 @@ def freeze_all_but_layernorms_(model):
     for m in model.modules():
         set_module_requires_grad_(m, isinstance(m, nn.LayerNorm))
 
+def unfreeze_last_n_layers_(enformer, n):
+    assert isinstance(enformer, Enformer)
+    transformer_blocks = enformer.transformer[1:]
+
+    for module in transformer_blocks[-n:]:
+        set_module_requires_grad_(module, True)
+
+def freeze_all_but_last_n_layers_(enformer, n):
+    assert isinstance(enformer, Enformer)
+    freeze_all_layers_(enformer)
+
+    transformer_blocks = enformer.transformer[1:]
+
+    for module in transformer_blocks[-n:]:
+        set_module_requires_grad_(module, False)
+
+# get enformer embeddings
+
 def get_enformer_embeddings(
     model,
     seq,

setup.py

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