clu_periodic_action_metric_writer_training_loop_quickstart.py

1import time
2from clu import periodic_actions
3from clu import metric_writers
4
5def train_loop(num_steps: int):
6  # 1. Initialize the metric writer (logs to TensorBoard by default)
7  writer = metric_writers.create_default_logdir("./logs")
8  
9  # 2. Setup periodic actions (e.g., logging every 10 steps or every 30 seconds)
10  report_progress = periodic_actions.ReportProgress(
11      num_train_steps=num_steps, writer=writer)
12  
13  for step in range(num_steps):
14    # --- Simulated training step start ---
15    time.sleep(0.1)
16    loss = 1.0 / (step + 1)
17    # --- Simulated training step end ---
18
19    # 3. Log metrics
20    if step % 10 == 0:
21      writer.write_scalars(step, {"loss": loss})
22    
23    # 4. Update periodic actions
24    report_progress(step)
25
26  # Ensure all metrics are flushed to disk
27  writer.flush()
28
29if __name__ == "__main__":
30  train_loop(num_steps=100)