record max step and min step (#341)

SwanHubX · Feb 22, 2024 · 6ec7f6d · 6ec7f6d
1 parent ce48cd8
commit 6ec7f6d
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/swanlab/data/run/exp.py b/swanlab/data/run/exp.py
@@ -195,8 +195,12 @@ def add(self, data: DataType, step: int = None):
         is_nan = self.__is_nan(data)
         if not is_nan:
             # 如果数据比之前的数据小，则更新最小值，否则不更新
-            self._summary["max"] = data if self._summary.get("max") is None else max(self._summary["max"], data)
-            self._summary["min"] = data if self._summary.get("min") is None else min(self._summary["min"], data)
+            if self._summary.get("max") is None or data > self._summary["max"]:
+                self._summary["max"] = data
+                self._summary["max_step"] = step
+            if self._summary.get("min") is None or data < self._summary["min"]:
+                self._summary["min"] = data
+                self._summary["min_step"] = step
         self._summary["num"] = self._summary.get("num", 0) + 1
         self.__steps.add(step)
         swanlog.debug(f"Add data, tag: {self.tag}, step: {step}, data: {data}")