iterations/n_iterations argument supported

src/maggotuba/models/modules.py

@@ -313,7 +313,8 @@ class DeepLinear(nn.Module):
 
 class MaggotClassifier(MaggotModule):
     def __init__(self, path, behavior_labels=[], n_latent_features=None,
-            n_layers=1, cfgfile=None, ptfile="trained_classifier.pt"):
+            n_layers=1, n_iterations=None, cfgfile=None,
+            ptfile="trained_classifier.pt"):
         super().__init__(path, cfgfile, ptfile)
         try: # try load config file, if any
             self.config
@@ -329,6 +330,16 @@ class MaggotClassifier(MaggotModule):
                 weight_init="xavier",
                 loss="cross-entropy",
                 optimizer="adam")
+            if n_iterations is not None:
+                if isinstance(n_iterations, str):
+                    n_iterations = map(int, n_iterations.split(','))
+                if isinstance(n_iterations, int):
+                    n_pretraining_iter = n_iterations // 2
+                    n_finetuning_iter = n_iterations // 2
+                else:
+                    n_pretraining_iter, n_finetuning_iter = n_iterations
+                self.config['pretraining_iter'] = n_pretraining_iter
+                self.config['finetuning_iter'] = n_finetuning_iter
 
     @classmethod
     def load_model(cls, config, path):
@@ -365,13 +376,22 @@ class MaggotClassifier(MaggotModule):
     def n_layers(self):
         return self.config["clf_depth"] + 1
 
+    @property
+    def n_pretraining_iter(self):
+        return self.config.get('pretraining_iter', None)
+
+    @property
+    def n_finetuning_iter(self):
+        return self.config.get('finetuning_iter', None)
+
 class SupervisedMaggot(nn.Module):
-    def __init__(self, cfgfilepath, behaviors=[], n_layers=1):
+    def __init__(self, cfgfilepath, behaviors=[], n_layers=1, n_epochs=None):
         super().__init__()
         if behaviors: # the model is only pre-trained
             self.encoder = PretrainedMaggotEncoder(cfgfilepath)
             self.clf = MaggotClassifier(self.encoder.path / "clf_config.json",
-                    behaviors, self.encoder.config["dim_latent"], n_layers)
+                    behaviors, self.encoder.config["dim_latent"], n_layers,
+                    n_epochs)
         else: # the model has been retrained
             self.clf = MaggotClassifier(cfgfilepath)
             self.encoder = MaggotEncoder(self.clf.config["autoencoder_config"],
@@ -398,15 +418,35 @@ class SupervisedMaggot(nn.Module):
         self.encoder.to(device)
         self.clf.to(device)
 
+    @property
+    def n_pretraining_iter(self):
+        n = self.clf.n_pretraining_iter
+        if n is None:
+            enc = self.encoder
+            n = enc.config['optim_iter']
+            if enc.was_pretrained():
+                n = n // 2
+        return n
+
+    @property
+    def n_finetuning_iter(self):
+        n = self.clf.n_finetuning_iter
+        if n is None:
+            enc = self.encoder
+            n = enc.config['optim_iter']
+            if enc.was_pretrained():
+                n = n // 2
+        return n
+
 class MultiscaleSupervisedMaggot(nn.Module):
-    def __init__(self, cfgfilepath, behaviors=[], n_layers=1):
+    def __init__(self, cfgfilepath, behaviors=[], n_layers=1, n_iterations=None):
         super().__init__()
         if behaviors: # the model is only pre-trained
             self.encoders = MaggotEncoders(cfgfilepath, cls=PretrainedMaggotEncoder)
             path = next(iter(self.encoders)).path.parent
             n_latent_features = sum(enc.config["dim_latent"] for enc in self.encoders)
             self.clf = MaggotClassifier(path / "clf_config.json",
-                    behaviors, n_latent_features, n_layers)
+                    behaviors, n_latent_features, n_layers, n_iterations)
         else: # the model has been retrained
             self.clf = MaggotClassifier(cfgfilepath)
             self.encoders = MaggotEncoders(self.clf.config["autoencoder_config"],
@@ -426,6 +466,26 @@ class MultiscaleSupervisedMaggot(nn.Module):
         self.clf.model # force parameter loading or initialization
         return super().parameters(self)
 
+    @property
+    def n_pretraining_iter(self):
+        n = self.clf.n_pretraining_iter
+        if n is None:
+            any_enc = self.encoders[0]
+            n = any_enc.config['optim_iter']
+            if any_enc.was_pretrained():
+                n = n // 2
+        return n
+
+    @property
+    def n_finetuning_iter(self):
+        n = self.clf.n_finetuning_iter
+        if n is None:
+            any_enc = self.encoders[0]
+            n = any_enc.config['optim_iter']
+            if any_enc.was_pretrained():
+                n = n // 2
+        return n
+
 """
 Bagging for `SupervisedMaggot`.
 
@@ -436,9 +496,10 @@ Bags of taggers are stored so that the models directory only contains
 subdirectories, each subdirectory specifying an individual tagger.
 """
 class MaggotBag(nn.Module):
-    def __init__(self, paths, behaviors=[], n_layers=1, cls=SupervisedMaggot):
+    def __init__(self, paths, behaviors=[], n_layers=1, n_iterations=None,
+                 cls=SupervisedMaggot):
         super().__init__()
-        self.maggots = [cls(path, behaviors, n_layers) for path in paths]
+        self.maggots = [cls(path, behaviors, n_layers, n_iterations) for path in paths]
         self._lead_maggot = None
 
     def forward(self, x):

src/maggotuba/models/train_model.py

@@ -6,7 +6,8 @@ import json
 import glob
 
 def train_model(backend, layers=1, pretrained_model_instance="default",
-                subsets=(1, 0, 0), rng_seed=None, balancing_strategy='maggotuba', **kwargs):
+                subsets=(1, 0, 0), rng_seed=None, iterations=1000,
+                balancing_strategy='maggotuba', **kwargs):
     # make_dataset generated or moved the larva_dataset file into data/interim/{instance}/
     #larva_dataset_file = backend.list_interim_files("larva_dataset_*.hdf5") # recursive
     larva_dataset_file = glob.glob(str(backend.interim_data_dir() / "larva_dataset_*.hdf5")) # not recursive (faster)
@@ -21,11 +22,11 @@ def train_model(backend, layers=1, pretrained_model_instance="default",
     # copy and load the pretrained model into the model instance directory
     if isinstance(pretrained_model_instance, str):
         config_file = import_pretrained_model(backend, pretrained_model_instance)
-        model = make_trainer(config_file, labels, layers)
+        model = make_trainer(config_file, labels, layers, iterations)
     else:
         pretrained_model_instances = pretrained_model_instance
         config_files = import_pretrained_models(backend, pretrained_model_instances)
-        model = make_trainer(config_files, labels, layers)
+        model = make_trainer(config_files, labels, layers, iterations)
     # fine-tune the model
     model.train(dataset)
     # add post-prediction rule ABC -> AAC

src/maggotuba/models/trainers.py

@@ -20,9 +20,9 @@ Training the model instead relies on the readily-preprocessed data of a
 *larva_dataset hdf5* file.
 """
 class MaggotTrainer:
-    def __init__(self, cfgfilepath, behaviors=[], n_layers=1,
+    def __init__(self, cfgfilepath, behaviors=[], n_layers=1, n_iterations=None,
             average_body_length=1.0, device=device):
-        self.model = SupervisedMaggot(cfgfilepath, behaviors, n_layers)
+        self.model = SupervisedMaggot(cfgfilepath, behaviors, n_layers, n_iterations)
         self.average_body_length = average_body_length # usually set later
         self.device = device
 
@@ -166,14 +166,12 @@ class MaggotTrainer:
         model.train() # this only sets the model in training mode (enables gradients)
         model.to(self.device)
         criterion = nn.CrossEntropyLoss(**kwargs)
-        nsteps = self.config['optim_iter']
         grad_clip = self.config['grad_clip']
         # pre-train the classifier with static encoder weights
         if model.encoder.was_pretrained():
-            nsteps = nsteps // 2
             optimizer = torch.optim.Adam(model.clf.parameters())
             print("pre-training the classifier...")
-            for step in range(nsteps):
+            for step in range(self.model.n_pretraining_iter):
                 optimizer.zero_grad()
                 # TODO: add an option for renormalizing the input
                 data, expected = self.draw(dataset)
@@ -186,7 +184,7 @@ class MaggotTrainer:
         optimizer = torch.optim.Adam(model.parameters())
         print(("fine-tuning" if model.encoder.was_pretrained() else "training") + \
                " the encoder and classifier...")
-        for step in range(nsteps):
+        for step in range(self.model.n_finetuning_iter):
             optimizer.zero_grad()
             data, expected = self.draw(dataset)
             predicted = self.forward(data, train=True)
@@ -246,20 +244,19 @@ def new_generator(seed=None):
 
 
 class MultiscaleMaggotTrainer(MaggotTrainer):
-    def __init__(self, cfgfilepath, behaviors=[], n_layers=1,
+    def __init__(self, cfgfilepath, behaviors=[], n_layers=1, n_iterations=None,
             average_body_length=1.0, device=device):
-        self.model = MultiscaleSupervisedMaggot(cfgfilepath, behaviors, n_layers)
+        self.model = MultiscaleSupervisedMaggot(cfgfilepath, behaviors,
+                                                n_layers, n_iterations)
         self.average_body_length = average_body_length # usually set later
         self.device = device
         self._default_encoder_config = None
         # check consistency
-        ref_config = self.config
-        for attr in ["batch_size", "optim_iter"]:
-            # TODO: add pretraining_iter and finetuning_iter parameters in
-            #       clf_config to have a lever other than optim_iter, that
-            #       could consequently be ignored
-            for enc in self.model.encoders:
-                assert enc.config[attr] == ref_config[attr]
+        if n_iterations is None:
+            ref_config = self.config
+            for attr in ["batch_size", "optim_iter"]:
+                for enc in self.model.encoders:
+                    assert enc.config[attr] == ref_config[attr]
 
     @property
     def config(self):