Add model prefix for loading

Author: MGlauer

chebai/cli.py

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

chebai/experiments.py

@@ -41,13 +41,14 @@ def model_kwargs(self, *args) -> Dict:
     def build_dataset(self, batch_size) -> datasets.XYBaseDataModule:
         raise NotImplementedError
 
-    def train(self, batch_size, *args):
+    def train(self, batch_size, *args, **kwargs):
         self.MODEL.run(
             self.dataset,
             self.MODEL.NAME,
             loss=self.LOSS,
             model_kwargs=self.model_kwargs(*args),
-            version=self.version
+            version=self.version,
+            **kwargs
         )
 
     def test(self, ckpt_path, *args):

chebai/models/base.py

@@ -146,7 +146,8 @@ def run(
         model_kwargs: dict = None,
         loss=torch.nn.BCELoss,
         weighted=False,
-        version=None
+        version=None,
+        **kwargs
     ):
         if model_args is None:
             model_args = []
@@ -199,7 +200,7 @@ def run(
 
         # Calculate weights per class
 
-        net = cls(*model_args, loss_cls=loss, **model_kwargs)
+        net = cls(*model_args, loss_cls=loss, **model_kwargs, **kwargs)
 
         # Early stopping seems to be bugged right now with ddp accelerator :(
         es = EarlyStopping(

chebai/models/electra.py

@@ -150,10 +150,15 @@ def __init__(self, **kwargs):
             kwargs["config"]["num_labels"] = self.out_dim
         self.config = ElectraConfig(**kwargs["config"], output_attentions=True)
         self.word_dropout = nn.Dropout(kwargs["config"].get("word_dropout", 0))
+        model_prefix = kwargs.get("load_prefix", None)
         if pretrained_checkpoint:
             with open(pretrained_checkpoint, "rb") as fin:
                 model_dict = torch.load(fin,map_location=self.device)
-                self.electra = ElectraModel.from_pretrained(None, state_dict=model_dict['state_dict'], config=self.config)
+                if model_prefix:
+                    state_dict = {str(k)[len(model_prefix):]:v for k,v in model_dict["state_dict"].items() if str(k).startswith(model_prefix)}
+                else:
+                    state_dict = model_dict["state_dict"]
+                self.electra = ElectraModel.from_pretrained(None, state_dict=state_dict, config=self.config)
         else:
             self.electra = ElectraModel(config=self.config)