allow larger amount of workers

Author: lfoppiano

delft/sequenceLabelling/trainer.py

@@ -30,7 +30,8 @@ def __init__(self,
                  save_path='',
                  preprocessor: Preprocessor=None,
                  transformer_preprocessor=None,
-                 enable_wandb = False
+                 enable_wandb = False,
+                 nb_workers=6
                  ):
 
         # for single model training
@@ -39,6 +40,8 @@ def __init__(self,
         # for n-folds training
         self.models = models
 
+        self.nb_workers = nb_workers
+
         self.embeddings = embeddings
         self.model_config = model_config
         self.training_config = training_config
@@ -192,7 +195,7 @@ def train_model(self, local_model, x_train, y_train, f_train=None,
                 model=local_model,
                 external_callbacks=callbacks
             )
-        nb_workers = 6
+        nb_workers = self.nb_workers
         multiprocessing = self.training_config.multiprocessing
 
         # multiple workers should work with transformer layers, but not with ELMo due to GPU memory limit (with GTX 1080Ti 11GB)

delft/sequenceLabelling/wrapper.py

@@ -1,3 +1,4 @@
+import multiprocessing
 import os
 
 from packaging import version
@@ -60,7 +61,7 @@
 class Sequence(object):
 
     # number of parallel worker for the data generator
-    nb_workers = 6
+    nb_workers = multiprocessing.cpu_count() - 1
 
     def __init__(
         self,
@@ -89,7 +90,8 @@ def __init__(
              multiprocessing=True,
              features_indices=None,
              transformer_name: str = None,
-             report_to_wandb = False
+             report_to_wandb = False,
+             nb_workers=6
          ):
 
         if model_name is None:
@@ -107,6 +109,7 @@ def __init__(
         self.embeddings_name = embeddings_name
 
         word_emb_size = 0
+        self.nb_workers = nb_workers
         self.embeddings = None
         self.model_local_path = None
 
@@ -127,22 +130,23 @@ def __init__(
             else:
                 learning_rate = 2e-5
 
-        self.model_config = ModelConfig(model_name=model_name,
-                                        architecture=architecture,
-                                        embeddings_name=embeddings_name,
-                                        word_embedding_size=word_emb_size,
-                                        char_emb_size=char_emb_size,
-                                        char_lstm_units=char_lstm_units,
-                                        max_char_length=max_char_length,
-                                        word_lstm_units=word_lstm_units,
-                                        max_sequence_length=max_sequence_length,
-                                        dropout=dropout,
-                                        recurrent_dropout=recurrent_dropout,
-                                        fold_number=fold_number,
-                                        batch_size=batch_size,
-                                        use_ELMo=use_ELMo,
-                                        features_indices=features_indices,
-                                        transformer_name=transformer_name)
+        self.model_config = ModelConfig(
+            model_name=model_name,
+            architecture=architecture,
+            embeddings_name=embeddings_name,
+            word_embedding_size=word_emb_size,
+            char_emb_size=char_emb_size,
+            char_lstm_units=char_lstm_units,
+            max_char_length=max_char_length,
+            word_lstm_units=word_lstm_units,
+            max_sequence_length=max_sequence_length,
+            dropout=dropout,
+            recurrent_dropout=recurrent_dropout,
+            fold_number=fold_number,
+            batch_size=batch_size,
+            use_ELMo=use_ELMo,
+            features_indices=features_indices,
+            transformer_name=transformer_name)
 
         self.training_config = TrainingConfig(learning_rate, batch_size, optimizer,
                                               lr_decay, clip_gradients, max_epoch,
@@ -259,7 +263,8 @@ def train_(self, x_train, y_train, f_train=None, x_valid=None, y_valid=None, f_v
             checkpoint_path=self.log_dir,
             preprocessor=self.p,
             transformer_preprocessor=self.model.transformer_preprocessor,
-            enable_wandb=self.report_to_wandb
+            enable_wandb=self.report_to_wandb,
+            nb_workers=self.nb_workers
         )
         trainer.train(x_train, y_train, x_valid, y_valid, features_train=f_train, features_valid=f_valid, callbacks=callbacks)
         if self.embeddings and self.embeddings.use_ELMo:
@@ -298,13 +303,16 @@ def train_nfold_(self, x_train, y_train, x_valid=None, y_valid=None, f_train=Non
             self.model_config.case_vocab_size = len(self.p.vocab_case)
             self.models = []
 
-        trainer = Trainer(self.model,
-                          self.models,
-                          self.embeddings,
-                          self.model_config,
-                          self.training_config,
-                          checkpoint_path=self.log_dir,
-                          preprocessor=self.p)
+        trainer = Trainer(
+            self.model,
+              self.models,
+              self.embeddings,
+              self.model_config,
+              self.training_config,
+              checkpoint_path=self.log_dir,
+              preprocessor=self.p,
+              nb_workers=self.nb_workers
+              )
 
         trainer.train_nfold(x_train, y_train, x_valid, y_valid, f_train=f_train, f_valid=f_valid, callbacks=callbacks)
         if self.embeddings and self.embeddings.use_ELMo: