Add epoch option

Author: MGlauer

chebai/cli.py

@@ -17,8 +17,9 @@ def cli():
 @click.option("-g", "--group", default="default")
 @click.option("--version", default=None)
 @click.option("--load-prefix", default=None)
+@click.option("--epochs", default=100)
 @click.argument("args", nargs=-1)
-def train(experiment, batch_size, group, version, load_prefix, args):
+def train(experiment, batch_size, group, version, load_prefix, epochs, args):
     """Run experiment identified by EXPERIMENT in batches of size BATCH_SIZE."""
     try:
         ex = experiments.EXPERIMENTS[experiment](batch_size, group, version=version)
@@ -27,7 +28,7 @@ def train(experiment, batch_size, group, version, load_prefix, args):
             "Experiment ID not found. The following are available:"
             + ", ".join(experiments.EXPERIMENTS.keys())
         )
-    ex.train(batch_size, *args, load_prefix=load_prefix)
+    ex.train(batch_size, epochs, *args, load_prefix=load_prefix)
 
 
 @click.command()

chebai/experiments.py

@@ -41,10 +41,11 @@ def model_kwargs(self, *args) -> Dict:
     def build_dataset(self, batch_size) -> datasets.XYBaseDataModule:
         raise NotImplementedError
 
-    def train(self, batch_size, *args, **kwargs):
+    def train(self, batch_size, epochs, *args, **kwargs):
         self.MODEL.run(
             self.dataset,
             self.MODEL.NAME,
+            epochs,
             loss=self.LOSS,
             model_kwargs=self.model_kwargs(*args),
             version=self.version,
@@ -93,7 +94,6 @@ def model_kwargs(self, *args) -> Dict:
                     type_vocab_size=1,
                 ),
             ),
-            epochs=100,
         )
 
     def build_dataset(self, batch_size) -> datasets.XYBaseDataModule:
@@ -117,7 +117,6 @@ def model_kwargs(self, *args) -> Dict:
                 num_hidden_layers=6,
                 type_vocab_size=1,
             ),
-            epochs=100,
         )
 
     def build_dataset(self, batch_size) -> datasets.XYBaseDataModule:
@@ -142,7 +141,6 @@ def model_kwargs(self, *args) -> Dict:
                 num_hidden_layers=6,
                 type_vocab_size=1,
             ),
-            epochs=100,
         )
 
     def build_dataset(self, batch_size) -> datasets.XYBaseDataModule:
@@ -174,8 +172,7 @@ def model_kwargs(self, *args) -> Dict:
                 num_attention_heads=8,
                 num_hidden_layers=6,
                 type_vocab_size=1,
-            ),
-            epochs=100,
+            )
         )
 
 
@@ -327,7 +324,6 @@ def model_kwargs(self, *args) -> Dict:
         return dict(
             in_length=50,
             hidden_length=100,
-            epochs=100,
         )
 
     def build_dataset(self, batch_size) -> datasets.XYBaseDataModule:
@@ -345,7 +341,6 @@ def model_kwargs(self, *args) -> Dict:
         return dict(
             in_length=50,
             hidden_length=100,
-            epochs=100,
         )
 
     def build_dataset(self, batch_size) -> datasets.XYBaseDataModule:

chebai/models/base.py

@@ -142,6 +142,7 @@ def run(
         cls,
         data,
         name,
+        epochs,
         model_args: list = None,
         model_kwargs: dict = None,
         loss=torch.nn.BCELoss,
@@ -209,7 +210,7 @@ def run(
 
         trainer = pl.Trainer(
             logger=tb_logger,
-            min_epochs=model_kwargs.get("epochs", 100),
+            min_epochs=epochs,
             callbacks=[best_checkpoint_callback, checkpoint_callback, es],
             replace_sampler_ddp=False,
             **trainer_kwargs