FairseqModelEnsemble class

pytext/task/tasks.py

@@ -71,7 +71,7 @@ class Config(NewTask.Config):
         ] = ClassificationMetricReporter.Config()
 
     def train_single_model(self, train_config, model_id, rank=0, world_size=1):
-        return self.trainer.train_single_model(
+        return self.trainer.real_trainers[model_id].train(
             self.data.batches(Stage.TRAIN),
             self.data.batches(Stage.EVAL),
             self.model.models[model_id],

pytext/trainers/ensemble_trainer.py

@@ -19,13 +19,18 @@ class Config(ConfigBase):
 
     @classmethod
     def from_config(cls, config: Config, model: torch.nn.Module, *args, **kwargs):
-        return cls(create_trainer(config.real_trainer, model, *args, **kwargs))
+        trainers = [
+            create_trainer(config.real_trainer, m, *args, **kwargs)
+            for m in model.models
+        ]
 
-    def __init__(self, real_trainer):
-        self.real_trainer = real_trainer
-        self.optimizer = real_trainer.optimizer
-        self.test = real_trainer.test
-        self.train_single_model = real_trainer.train
+        return cls(trainers)
+
+    def __init__(self, real_trainers):
+        self.real_trainers = real_trainers
+        self.optimizer = real_trainers[0].optimizer
+        self.test = real_trainers[0].test
+        self.train_single_model = real_trainers[0].train
 
     """
     Train and eval ensemble model, each sub model will be trained separately in
@@ -53,7 +58,7 @@ def __init__(self, real_trainer):
     """
 
     def train(self, train_iter, eval_iter, model, *args, **kwargs):
-        for m in model.models:
-            self.train_single_model(train_iter, eval_iter, m, *args, **kwargs)
+        for (m, t) in zip(model.models, self.real_trainers):
+            t.train(train_iter, eval_iter, m, *args, **kwargs)
         model.merge_sub_models()
         return model, None