Kalash x transformers

ptls/frames/coles/__init__.py

@@ -1,5 +1,5 @@
 from .coles_dataset import ColesDataset, ColesIterableDataset
 from .coles_supervised_dataset import ColesSupervisedDataset, ColesSupervisedIterableDataset
-from .coles_module import CoLESModule
+from .coles_module import CoLESModule, CoLESModuleWarmup
 from .coles_supervised_module import ColesSupervisedModule
 

ptls/frames/coles/coles_module.py

@@ -74,4 +74,41 @@ def shared_step(self, x, y):
         y_h = self(x)
         if self._head is not None:
             y_h = self._head(y_h)
-        return y_h, y
+        return y_h, y
+
+class CoLESModuleWarmup(CoLESModule):
+    def __init__(self,
+                 seq_encoder: SeqEncoderContainer = None,
+                 head=None,
+                 loss=None,
+                 validation_metric=None,
+                 optimizer_partial=None,
+                 lr_scheduler_partial=None,
+                 warmup_steps = 500,
+                 initial_lr = 0.001):
+
+        super().__init__(seq_encoder,
+                         head,
+                         loss,
+                         validation_metric,
+                         optimizer_partial,
+                         lr_scheduler_partial)
+        self.warmup_steps = warmup_steps
+        self.initial_lr = initial_lr
+
+    def optimizer_step(self, 
+                       epoch, 
+                       batch_idx, 
+                       optimizer, 
+                       optimizer_idx, 
+                       optimizer_closure, 
+                       on_tpu=False, 
+                       using_native_amp=False, 
+                       using_lbfgs=False):
+
+        optimizer.step(closure = optimizer_closure)
+        if self.trainer.global_step < self.warmup_steps:
+            lr_scale = min(1., float(self.trainer.global_step + 1) / self.warmup_steps)
+            for pg in optimizer.param_groups:
+                pg['lr'] = lr_scale * self.initial_lr
+

ptls/nn/__init__.py

@@ -5,7 +5,7 @@
 from .seq_encoder import (
     RnnEncoder,  TransformerEncoder, LongformerEncoder,
     RnnSeqEncoder, TransformerSeqEncoder, LongformerSeqEncoder, AggFeatureSeqEncoder,
-    GptEncoder
+    GptEncoder, XTransformerEncoder, XTransformerSeqEncoder
 )
 
 from .pb import PBDropout, PBLinear, PBL2Norm, PBLayerNorm, PBReLU

ptls/nn/seq_encoder/__init__.py

@@ -3,6 +3,7 @@
 from .longformer_encoder import LongformerEncoder
 from .gpt_encoder import GptEncoder
 from .custom_encoder import Encoder
+from .x_transformer import XTransformerEncoder, XTransformerSeqEncoder
 
 from .containers import RnnSeqEncoder, TransformerSeqEncoder, LongformerSeqEncoder, CustomSeqEncoder
 from .agg_feature_seq_encoder import AggFeatureSeqEncoder