add setting for tqdm mininterval

Author: vpratz

bayesflow/default_settings.py

@@ -199,3 +199,8 @@ def __init__(self, meta_dict: dict, mandatory_fields: list = []):
 
 
 MMD_BANDWIDTH_LIST = [1e-6, 1e-5, 1e-4, 1e-3, 1e-2, 1e-1, 1, 5, 10, 15, 20, 25, 30, 35, 100, 1e3, 1e4, 1e5, 1e6]
+
+# Minimum time interval between tqdm status updates to reduce
+# load. Only respected when refresh=False in set_postfix
+# and set_postfix_str
+TQDM_MININTERVAL = 0.1  # in seconds

bayesflow/simulation.py

@@ -29,7 +29,7 @@
 
 logging.basicConfig()
 
-from bayesflow.default_settings import DEFAULT_KEYS
+from bayesflow.default_settings import DEFAULT_KEYS, TQDM_MININTERVAL
 from bayesflow.diagnostics import plot_prior2d
 from bayesflow.exceptions import ConfigurationError
 
@@ -1109,7 +1109,9 @@ def presimulate_and_save(
         # Generate the presimulation files
         file_counter = extend_from
         for i in range(total_files):
-            with tqdm(total=batches_per_file, desc=f"Batches generated for file {i+1}") as p_bar:
+            with tqdm(
+                total=batches_per_file, desc=f"Batches generated for file {i+1}", mininterval=TQDM_MININTERVAL
+            ) as p_bar:
                 file_list = [{} for _ in range(batches_per_file)]
                 for k in range(batches_per_file):
                     file_list[k] = self.__call__(batch_size=batch_size)

bayesflow/trainers.py

@@ -34,7 +34,7 @@
 )
 from bayesflow.computational_utilities import maximum_mean_discrepancy
 from bayesflow.configuration import *
-from bayesflow.default_settings import DEFAULT_KEYS, OPTIMIZER_DEFAULTS
+from bayesflow.default_settings import DEFAULT_KEYS, OPTIMIZER_DEFAULTS, TQDM_MININTERVAL
 from bayesflow.diagnostics import plot_latent_space_2d, plot_sbc_histograms
 from bayesflow.exceptions import ArgumentError, SimulationError
 from bayesflow.helper_classes import (
@@ -432,7 +432,7 @@ def train_online(
 
         # Loop through training epochs
         for ep in range(1, epochs + 1):
-            with tqdm(total=iterations_per_epoch, desc=f"Training epoch {ep}") as p_bar:
+            with tqdm(total=iterations_per_epoch, desc=f"Training epoch {ep}", mininterval=TQDM_MININTERVAL) as p_bar:
                 for it in range(1, iterations_per_epoch + 1):
                     # Perform one training step and obtain current loss value
                     loss = self._train_step(batch_size, update_step=_backprop_step, **kwargs)
@@ -450,7 +450,7 @@ def train_online(
                     disp_str = format_loss_string(ep, it, loss, avg_dict, lr=lr)
 
                     # Update progress bar
-                    p_bar.set_postfix_str(disp_str)
+                    p_bar.set_postfix_str(disp_str, refresh=False)
                     p_bar.update(1)
 
             # Store and compute validation loss, if specified
@@ -558,7 +558,9 @@ def train_offline(
 
         # Loop through epochs
         for ep in range(1, epochs + 1):
-            with tqdm(total=data_set.num_batches, desc="Training epoch {}".format(ep)) as p_bar:
+            with tqdm(
+                total=data_set.num_batches, desc="Training epoch {}".format(ep), mininterval=TQDM_MININTERVAL
+            ) as p_bar:
                 # Loop through dataset
                 for bi, forward_dict in enumerate(data_set, start=1):
                     # Perform one training step and obtain current loss value
@@ -578,7 +580,7 @@ def train_offline(
                     disp_str = format_loss_string(ep, bi, loss, avg_dict, lr=lr, it_str="Batch")
 
                     # Update progress
-                    p_bar.set_postfix_str(disp_str)
+                    p_bar.set_postfix_str(disp_str, refresh=False)
                     p_bar.update(1)
 
             # Store and compute validation loss, if specified
@@ -731,7 +733,7 @@ def train_from_presimulation(
                     f"Loading a simulation file resulted in a {type(epoch_data)}. Must be a dictionary or a list!"
                 )
 
-            with tqdm(total=len(index_list), desc=f"Training epoch {ep}") as p_bar:
+            with tqdm(total=len(index_list), desc=f"Training epoch {ep}", mininterval=TQDM_MININTERVAL) as p_bar:
                 for it, index in enumerate(index_list, start=1):
                     # Perform one training step and obtain current loss value
                     input_dict = self.configurator(epoch_data[index])
@@ -756,7 +758,7 @@ def train_from_presimulation(
                     disp_str = format_loss_string(ep, it, loss, avg_dict, lr=lr)
 
                     # Update progress bar
-                    p_bar.set_postfix_str(disp_str)
+                    p_bar.set_postfix_str(disp_str, refresh=False)
                     p_bar.update(1)
 
             # Store after each epoch, if specified
@@ -873,7 +875,7 @@ def train_experience_replay(
 
         # Loop through epochs
         for ep in range(1, epochs + 1):
-            with tqdm(total=iterations_per_epoch, desc=f"Training epoch {ep}") as p_bar:
+            with tqdm(total=iterations_per_epoch, desc=f"Training epoch {ep}", mininterval=TQDM_MININTERVAL) as p_bar:
                 for it in range(1, iterations_per_epoch + 1):
                     # Simulate a batch of data and store into buffer
                     input_dict = self._forward_inference(
@@ -900,7 +902,7 @@ def train_experience_replay(
                     disp_str = format_loss_string(ep, it, loss, avg_dict, lr=lr)
 
                     # Update progress bar
-                    p_bar.set_postfix_str(disp_str)
+                    p_bar.set_postfix_str(disp_str, refresh=False)
                     p_bar.update(1)
 
             # Store and compute validation loss, if specified
@@ -1099,7 +1101,7 @@ def mmd_hypothesis_test(
         num_reference = reference_summary.shape[0]
 
         mmd_null_samples = np.empty(num_null_samples, dtype=np.float32)
-        for i in tqdm(range(num_null_samples)):
+        for i in tqdm(range(num_null_samples), mininterval=TQDM_MININTERVAL):
             if bootstrap:
                 bootstrap_idx = np.random.randint(0, num_reference, size=num_observed)
                 simulated_summary = tf.gather(reference_summary, bootstrap_idx, axis=0)