Clear metric nonfinite

lightvector · Mar 10, 2024 · aea2420 · aea2420
1 parent 44b5447
commit aea2420
Showing 1 changed file with 10 additions and 0 deletions.
diff --git a/python/train.py b/python/train.py
@@ -909,6 +909,14 @@ def log_metrics(metric_sums, metric_weights, metrics, metrics_out):
             metrics_out.write(json.dumps(metrics_to_print) + "\n")
             metrics_out.flush()
 
+    def clear_metric_nonfinite(metric_sums, metric_weights):
+        for metric in metric_sums:
+            if not math.isfinite(metric_sums[metric]):
+                logging.warning(f"NONFINITE VALUE OF METRIC {metric}, CLEARING IT BACK TO EMPTY")
+                metric_sums[metric] = 0.0
+                metric_weights[metric] = 0.0
+
+
     if rank == 0:
         train_metrics_out = open(os.path.join(traindir,"metrics_train.json"),"a")
         val_metrics_out = open(os.path.join(traindir,"metrics_val.json"),"a")
@@ -975,6 +983,8 @@ def log_metrics(metric_sums, metric_weights, metrics, metrics_out):
         logging.info("GC collect")
         gc.collect()
 
+        clear_metric_nonfinite(running_metrics["sums"], running_metrics["weights"])
+
         logging.info("=========================================================================")
         logging.info("BEGINNING NEXT EPOCH " + str(num_epochs_this_instance))
         logging.info("=========================================================================")