File size: 2,431 Bytes
7cffee1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 | import argparse
import csv
from collections import defaultdict
from itertools import cycle
from pathlib import Path
import matplotlib.pyplot as plt
COLORS = ["#0072B2", "#D55E00", "#009E73"]
def read_metrics(path: Path) -> dict[str, list[dict[str, float]]]:
metrics = defaultdict(list)
with path.open(encoding="utf-8", newline="") as metrics_file:
for row in csv.DictReader(metrics_file):
metrics[row["run"]].append(
{
"tokens": float(row["tokens"]),
"loss": float(row["loss"]),
"wall_time_seconds": float(row["wall_time_seconds"]),
}
)
if not metrics:
raise ValueError(f"No metric rows found in {path}")
return dict(metrics)
def plot_metrics(
metrics: dict[str, list[dict[str, float]]],
x_field: str,
x_scale: float,
x_label: str,
output_path: Path,
) -> None:
figure, axis = plt.subplots(figsize=(10, 6))
for color, (run, rows) in zip(cycle(COLORS), metrics.items()):
x_values = [row[x_field] / x_scale for row in rows]
loss_values = [row["loss"] for row in rows]
axis.plot(x_values, loss_values, color=color, linewidth=1.8, label=run)
axis.set_xlabel(x_label)
axis.set_ylabel("Training loss")
axis.set_title(f"Training loss vs {x_label.lower()}")
axis.grid(True, alpha=0.25)
axis.legend(title="Run", fontsize=8)
figure.tight_layout()
figure.savefig(output_path, dpi=160)
plt.close(figure)
def main() -> None:
argument_parser = argparse.ArgumentParser()
argument_parser.add_argument("metrics", type=Path)
argument_parser.add_argument("--output-dir", type=Path, required=True)
args = argument_parser.parse_args()
metrics = read_metrics(args.metrics)
args.output_dir.mkdir(parents=True, exist_ok=True)
plot_metrics(
metrics,
x_field="wall_time_seconds",
x_scale=3600,
x_label="Wall-clock time (hours)",
output_path=args.output_dir / "loss_vs_time.png",
)
plot_metrics(
metrics,
x_field="tokens",
x_scale=1_000_000,
x_label="Tokens processed (millions)",
output_path=args.output_dir / "loss_vs_tokens.png",
)
print(f"Wrote {args.output_dir / 'loss_vs_time.png'}")
print(f"Wrote {args.output_dir / 'loss_vs_tokens.png'}")
if __name__ == "__main__":
main()
|