LSTM optim optuna

Author: Wonters

src/mixins.py

@@ -38,28 +38,26 @@ def optuna_train(self, run_name:str = "", n_trials:int=30, frac=0.1):
         study = optuna.create_study(direction="maximize",
                                     pruner=optuna.pruners.MedianPruner(n_startup_trials=5, n_warmup_steps=1))
         study.optimize(self.objective, n_trials=n_trials)
-        
-
-    def objective(self, trial):
+    
+    def params_optim(self, trial):
         lr = trial.suggest_loguniform('lr', 1e-6, 1e-3)
         gamma = trial.suggest_float('gamma', 0.1, 0.9)
         step_size = trial.suggest_int('step_size', 2, 10)
+        return {'lr': lr, 'gamma': gamma, 'step_size': step_size} 
+
+    def objective(self, trial):
+        kwargs = self.params_optim(trial)
         with mlflow.start_run(nested=True):
-            mlflow.log_params({
-                "lr": lr,
-                "gamma": gamma,
-                "step_size": step_size
-            })
-            self.optimizer = torch.optim.Adam(self.model.parameters(), lr=lr)
-            self.scheduler = torch.optim.lr_scheduler.StepLR(self.optimizer, step_size=step_size, gamma=gamma)
+            mlflow.log_params(kwargs)
+            self.reinit_scheduler_optimizer(**kwargs)
             acc = self.train()
         return acc
 
     def _train_batch(self, x, y):
         inputs = self.tokenizer(x, return_tensors="pt", truncation=True, padding=True)
         if isinstance(inputs, dict) and inputs["input_ids"]:
             inputs["input_ids"] = inputs["input_ids"].float()
-        if isinstance(y, torch.Tensor) and y.dtype == torch.float32:
+        if False and isinstance(y, torch.Tensor) and y.dtype == torch.float32:
             labels = y.long()
         else:
             labels = y.float()
@@ -69,9 +67,10 @@ def _train_batch(self, x, y):
         outputs = self.model(**inputs)
         try:
             loss = self.criterion(outputs.logits, labels)
+            _, preds = torch.max(outputs.logits, dim=1)
         except AttributeError:
             loss = self.criterion(outputs, labels)
-        _, preds = torch.max(outputs.logits, dim=1)
+            preds = outputs 
         correct = (preds == labels).sum().item()    
         acc = correct / len(labels)
         loss.backward()
@@ -99,8 +98,10 @@ def train(self):
         for epoch in tqdm(range(self.epoch)):
             for tweets, labels in tqdm(self.dataloader):
                 try:
-                    acc.append(self._train_batch(tweets, labels.float()))
+                    current_acc = self._train_batch(tweets, labels.float())
+                    acc.append(current_acc)
                 except RuntimeError as e:
+                    raise e
                     logger.error(e)
                     del tweets, labels, self.optimizer
                     gc.collect()
@@ -120,7 +121,10 @@ def train(self):
                     logger.info(
                         f"CUDA allocated memory: {torch.cuda.memory_allocated()}"
                     )
-            self.scheduler.step()
+            if isinstance(self.scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau):
+                self.scheduler.step(current_acc)
+            else:   
+                self.scheduler.step()
         super().train()
         return sum(acc)/len(acc)
 

src/ml.py

@@ -357,6 +357,16 @@ class BertModel(TorchBaseModel):
     lr = 2.561e-4
     device = DEVICE
 
+    def params_optim(self, trial):
+        lr = trial.suggest_loguniform('lr', 1e-6, 1e-3)
+        gamma = trial.suggest_float('gamma', 0.1, 0.9)
+        step_size = trial.suggest_int('step_size', 2, 10)
+        return {'lr': lr, 'gamma': gamma, 'step_size': step_size} 
+
+    def reinit_scheduler_optimizer(self, lr, gamma, step_size):
+        self.optimizer = torch.optim.Adam(self.model.parameters(), lr=lr)
+        self.scheduler = torch.optim.lr_scheduler.StepLR(self.optimizer, step_size=step_size, gamma=gamma)
+
     def load_checkpoint(self):
         if Path(self.checkpoint).exists():
             self.tokenizer = AutoTokenizer.from_pretrained(self.checkpoint)
@@ -449,27 +459,40 @@ def load_checkpoint(self):
             self.model.load_state_dict(embedding_weights, strict=False)
             self.model.eval()
         self.optimizer = torch.optim.Adam(self.model.parameters(), lr=self.lr)
-        # self.scheduler = torch.optim.lr_scheduler.StepLR(self.optimizer, step_size=5, gamma=0.1)
         self.scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
             self.optimizer, mode="min", factor=0.5, patience=2
         )
         self.criterion = torch.nn.BCEWithLogitsLoss()
 
+    def params_optim(self, trial):
+        lr = trial.suggest_loguniform('lr', 1e-6, 1e-3)
+        factor = trial.suggest_float('factor', 0.1, 0.9)
+        patience = trial.suggest_int('patience', 2, 10)
+        return {'lr': lr, 'factor': factor, 'patience': patience} 
+
+    def reinit_scheduler_optimizer(self, lr, factor, patience):
+        self.optimizer = torch.optim.Adam(self.model.parameters(), lr=lr)
+        self.scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
+            self.optimizer, mode="min", factor=factor, patience=patience
+        )
+
     def save(self):
         # Create the parent directory saving tokenizer
         self.tokenizer.save_pretrained(self.checkpoint)
         torch.save(self.model.state_dict(), self.checkpoint + "/model.pth")
 
     def predict(self, x):
-        inputs = self.preprocessing(x)
-        inputs = inputs.to(self.device)
-        with torch.no_grad():
-            outputs = self.model(**inputs)
-            # Appliquer sigmoïde sur les 4 prédictions
-            probs = torch.sigmoid(outputs)
-            # Convertir en classes (0 ou 1) en utilisant un seuil de 0.5
-            predicted_classes = (probs > 0.5).int()
-        return predicted_classes.tolist()
+        predicted_class = []
+        for i in range(0, len(x), self.batch_size):
+            inputs = self.preprocessing(x[i:i+self.batch_size])
+            inputs = inputs.to(self.device)
+            with torch.no_grad():
+                outputs = self.model(**inputs)
+                # Appliquer sigmoïde sur les 4 prédictions
+                probs = torch.sigmoid(outputs)
+                # Convertir en classes (0 ou 1) en utilisant un seuil de 0.5
+                predicted_class.extend((probs > 0.5).int().tolist())
+        return predicted_class
 
 
 def split_data(df: pd.DataFrame, shuffle: bool = True):

train.py

@@ -8,4 +8,4 @@
 file = "../data/training.1600000.processed.noemoticon.csv"
 original_df = load_data(file)
 model = LSTMModel(original_df)
-model.optuna_train(n_trials=5, frac=0.01)
+model.optuna_train(n_trials=5, frac=0.001)