From 6a5b139f3771201ca1e70e6da1531d66b1dd265c Mon Sep 17 00:00:00 2001
From: ppdebreuck <pierre-paul.debreuck@uclouvain.be>
Date: Sat, 13 Nov 2021 16:44:03 +0100
Subject: [PATCH] reduc tf logging when doing hyper search

---
 modnet/hyper_opt/fit_genetic.py | 7 +++++++
 modnet/models/ensemble.py       | 7 +++++++
 modnet/models/vanilla.py        | 7 +++++++
 3 files changed, 21 insertions(+)

diff --git a/modnet/hyper_opt/fit_genetic.py b/modnet/hyper_opt/fit_genetic.py
index f4a2807e..ea2c8250 100644
--- a/modnet/hyper_opt/fit_genetic.py
+++ b/modnet/hyper_opt/fit_genetic.py
@@ -308,6 +308,11 @@ def function_fitness(
         """
 
         from modnet.matbench.benchmark import matbench_kfold_splits
+        import os
+
+        os.environ[
+            "TF_CPP_MIN_LOG_LEVEL"
+        ] = "2"  # many models will be fitted => reduce output
 
         num_nested_folds = 5
         if nested:
@@ -381,6 +386,8 @@ def function_fitness(
         pool.close()
         pool.join()
 
+        os.environ["TF_CPP_MIN_LOG_LEVEL"] = "0"  # reset
+
         return val_loss_per_individual, np.array(models), np.array(individuals)
 
     def run(
diff --git a/modnet/models/ensemble.py b/modnet/models/ensemble.py
index 52fde07d..d03a7195 100644
--- a/modnet/models/ensemble.py
+++ b/modnet/models/ensemble.py
@@ -244,6 +244,11 @@ def fit_preset(
         """
 
         from modnet.matbench.benchmark import matbench_kfold_splits
+        import os
+
+        os.environ[
+            "TF_CPP_MIN_LOG_LEVEL"
+        ] = "2"  # many models will be fitted => reduce output
 
         if callbacks is None:
             es = tf.keras.callbacks.EarlyStopping(
@@ -403,6 +408,8 @@ def fit_preset(
                     final_models += models[idx][i].model
             self.__init__(modnet_models=final_models)
 
+        os.environ["TF_CPP_MIN_LOG_LEVEL"] = "0"  # reset
+
         return models, val_losses, best_learning_curve, learning_curves, best_preset
 
     def _make_picklable(self):
diff --git a/modnet/models/vanilla.py b/modnet/models/vanilla.py
index a725eb6d..1aac0faf 100644
--- a/modnet/models/vanilla.py
+++ b/modnet/models/vanilla.py
@@ -411,6 +411,11 @@ def fit_preset(
         """
 
         from modnet.matbench.benchmark import matbench_kfold_splits
+        import os
+
+        os.environ[
+            "TF_CPP_MIN_LOG_LEVEL"
+        ] = "2"  # many models will be fitted => reduce output
 
         if callbacks is None:
             es = tf.keras.callbacks.EarlyStopping(
@@ -559,6 +564,8 @@ def fit_preset(
             self.model = best_model.model
             self._scaler = best_model._scaler
 
+        os.environ["TF_CPP_MIN_LOG_LEVEL"] = "0"  # reset
+
         return models, val_losses, best_learning_curve, learning_curves, best_preset
 
     def predict(self, test_data: MODData, return_prob=False) -> pd.DataFrame: