add checkpointing to transformers trunk, to save on memory when finetuning

Author: lucidrains

enformer_pytorch/enformer_pytorch.py

@@ -1,6 +1,7 @@
 import math
 import torch
 from torch import nn, einsum
+from torch.utils.checkpoint import checkpoint_sequential
 from einops import rearrange, reduce
 from einops.layers.torch import Rearrange
 
@@ -262,7 +263,8 @@ def __init__(
         attn_dim_key = 64,
         dropout_rate = 0.4,
         attn_dropout = 0.05,
-        pos_dropout = 0.01
+        pos_dropout = 0.01,
+        use_checkpointing = False
     ):
         super().__init__()
         self.dim = dim
@@ -359,6 +361,10 @@ def __init__(
             nn.Softplus()
         ), output_heads))
 
+        # use checkpointing on transformer trunk
+
+        self.use_checkpointing = use_checkpointing
+
     def set_target_length(self, target_length):
         crop_module = self._trunk[-2]
         crop_module.target_length = target_length
@@ -370,7 +376,20 @@ def trunk(self):
     @property
     def heads(self):
         return self._heads
-    
+
+    def trunk_checkpointed(self, x):
+        x = self.stem(x)
+        x = self.conv_tower(x)
+        x = self.transformer[0](x)
+
+        # todo (move the rearrange out of self.transformers sequential module, and transfer all weights to new module rearrangement, directly checkpoint on self.transformers)
+        transformer_blocks = self.transformer[1:]
+        x = checkpoint_sequential(nn.Sequential(*transformer_blocks), len(transformer_blocks), x)
+
+        x = self.crop_final(x)
+        x = self.final_pointwise(x)
+        return x
+
     def forward(
         self,
         x,
@@ -390,7 +409,8 @@ def forward(
         if no_batch:
             x = rearrange(x, '... -> () ...')
 
-        x = self._trunk(x)
+        trunk_fn = self.trunk_checkpointed if self.use_checkpointing else self._trunk
+        x = trunk_fn(x)
 
         if no_batch:
             x = rearrange(x, '() ... -> ...')

setup.py

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