Logger for saving metrics to a local file, GCS and TensorBoard.
| 85 | |
| 86 | |
| 87 | class MetricLogger: |
| 88 | """ |
| 89 | Logger for saving metrics to a local file, GCS and TensorBoard. |
| 90 | """ |
| 91 | |
| 92 | def __init__(self, config, learning_rate_schedule): |
| 93 | self.writer = max_utils.initialize_summary_writer(config.tensorboard_dir, config.run_name) |
| 94 | self.config = config |
| 95 | self.metadata = {} |
| 96 | self.running_gcs_metrics = [] if config.gcs_metrics else None |
| 97 | self.performance_metric_queue = self.get_performance_metric_queue(config) |
| 98 | self.learning_rate_schedule = learning_rate_schedule |
| 99 | self.cumulative_eval_metrics = {"scalar": defaultdict(float)} |
| 100 | self.buffered_train_metrics = None |
| 101 | if self.config.managed_mldiagnostics: |
| 102 | ManagedMLDiagnostics(config) # Initialize the MLRun instance. |
| 103 | |
| 104 | def reset_eval_metrics(self): |
| 105 | """Resets the cumulative metrics dictionary for a new evaluation run.""" |
| 106 | self.cumulative_eval_metrics = {"scalar": defaultdict(float)} |
| 107 | |
| 108 | def write_metrics(self, metrics, step, is_training=True): |
| 109 | """Entry point for all metrics writing in Train's Main.""" |
| 110 | if metrics: |
| 111 | self.log_metrics(metrics, step, is_training) |
| 112 | |
| 113 | if self.config.enable_tensorboard: |
| 114 | self.write_metrics_to_tensorboard(metrics, step, is_training) |
| 115 | |
| 116 | if self.config.metrics_file: |
| 117 | self.write_metrics_locally(metrics, step) |
| 118 | |
| 119 | if self.config.gcs_metrics and jax.process_index() == 0: |
| 120 | self.write_metrics_for_gcs(metrics, step, is_training) |
| 121 | |
| 122 | if self.config.managed_mldiagnostics: |
| 123 | self.write_metrics_to_managed_mldiagnostics(metrics, step) |
| 124 | |
| 125 | def log_metrics(self, metrics, step, is_training): |
| 126 | """Logs metrics via max_logging.""" |
| 127 | if is_training: |
| 128 | self._log_training_metrics(metrics, step) |
| 129 | else: |
| 130 | self._log_eval_metrics(metrics, step) |
| 131 | |
| 132 | def _log_training_metrics(self, metrics, step): |
| 133 | """Handles training-specific metric logging.""" |
| 134 | # Skip logging if in profiler activation/deactivation steps |
| 135 | # TODO(b/456828037): Switch to subprocess profiling to avoid timing artifacts at boundary steps. |
| 136 | scalars = metrics["scalar"] |
| 137 | loss = scalars["learning/loss"] |
| 138 | is_rampup = step < self.config.rampup_end_step |
| 139 | is_metric_hidden_step = self.config.hide_profiler_step_metric and self._is_profiler_boundary_step(step) |
| 140 | |
| 141 | # Start building the log parts |
| 142 | log_parts = [] |
| 143 | if is_rampup: |
| 144 | log_parts.append("[Rampup Batch Size Phase]") |
no outgoing calls
no test coverage detected