Transformer architecture implemented in TypeScript with WebGPU.
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 
 
nn/tools/stats.py

178 lines
6.0 KiB

import sys
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.animation as animation
import matplotlib.dates as mdates # Add this import at the top
import select
from datetime import datetime
import threading
import queue
fig = plt.figure(figsize=(10, 6))
ax_lr = fig.add_axes([0.05, 0.3, 0.9, 0.65])
ax_lr.set_facecolor("none")
ax_lr.set_ylim(bottom=0, top=None)
ax_lr.yaxis.set_ticks_position('right')
ax_lr.tick_params(axis='y', labelcolor='tab:brown')
ax_acc = fig.add_axes([0.05, 0.3, 0.9, 0.65], sharex=ax_lr)
ax_acc.set_facecolor("none")
ax_acc.set_ylim(0, 100)
ax_acc.yaxis.set_ticks_position('right')
ax_acc.tick_params(axis='y', labelcolor='tab:green')
ax_loss = fig.add_axes([0.05, 0.3, 0.9, 0.65], sharex=ax_lr)
ax_loss.set_facecolor("none")
ax_loss.set_ylim(bottom=0, top=None)
ax_loss.tick_params(axis='y', labelcolor='tab:blue')
ax_time = fig.add_axes([0.05, 0.05, 0.9, 0.2])
ax_time.set_facecolor("none")
ax_time.set_ylim(bottom=0, top=None)
ax_time.yaxis.set_ticks_position('right')
ax_time.tick_params(axis='y', labelcolor='tab:orange')
columns = ["index", "timestamp", "epoch", "batch", "lr", "loss", "accuracy", "grad_inp", "grad_qry", "grad_key", "grad_val", "grad_out", "queue_ms", "compute_ms", "infer_ms"]
df = pd.DataFrame(columns=columns)
ema_alpha = 0.1
window_size = 50
def parse_line(line):
try:
parts = line.strip().split(",")
if parts[0] == "index" or len(parts) < len(columns):
return None
row = {
"index": int(parts[0]),
"timestamp": datetime.fromtimestamp(int(parts[1]) / 1000),
"epoch": int(parts[2]),
"batch": int(parts[3]),
"lr": float(parts[4]),
"loss": float(parts[5]),
"accuracy": float(parts[6]) * 100,
"grad_inp": float(parts[7]),
"grad_qry": float(parts[8]),
"grad_key": float(parts[9]),
"grad_val": float(parts[10]),
"grad_out": float(parts[11]),
"queue_ms": float(parts[12]),
"compute_ms": float(parts[13]),
"infer_ms": float(parts[14]),
}
return row
except Exception as e:
print(f"Bad line: {line.strip()} ({e})", file=sys.stderr)
return None
line_queue = queue.Queue()
def stdin_reader():
for line in sys.stdin:
line_queue.put(line)
threading.Thread(target=stdin_reader, daemon=True).start()
def get_new_lines():
lines = []
while not line_queue.empty():
lines.append(line_queue.get())
return lines
last_seen_index = None
def animate(_):
global df
lines = get_new_lines()
if not lines:
return
new_rows = [parse_line(line) for line in lines]
new_rows = [r for r in new_rows if r is not None]
if not new_rows:
return
global last_seen_index
new_index = new_rows[-1]["index"]
if last_seen_index and (new_index < last_seen_index):
print("Index reset, restarting plot")
df = pd.DataFrame(columns=columns)
last_seen_index = new_index
df_new = pd.DataFrame(new_rows)
df_new["total_ms"] = df_new["queue_ms"] + df_new["compute_ms"] + df_new["infer_ms"]
df = pd.concat([df, df_new], ignore_index=True)
df["ema_loss"] = df["loss"].ewm(alpha=ema_alpha).mean().combine_first(df["loss"])
df["avg_loss"] = df["loss"].rolling(window=window_size).mean().combine_first(df["loss"])
df["avg_accuracy"] = df["accuracy"].rolling(window=window_size).mean().combine_first(df["accuracy"])
df["avg_total_ms"] = df["total_ms"].rolling(window=window_size).mean().combine_first(df["total_ms"])
ax_lr.clear()
ax_lr.set_ylabel("LR")
ax_lr.yaxis.set_label_position("right")
ax_acc.clear()
ax_acc.set_ylabel("Accuracy (%)")
ax_loss.clear()
ax_loss.set_title("Live Training Metrics")
ax_loss.set_xlabel("Epoch / Batch")
ax_loss.set_ylabel("Loss")
ax_acc.yaxis.set_label_position("right")
ax_time.clear()
ax_time.set_xlabel("Time")
ax_time.set_ylabel("Processing Time (ms)")
ax_time.yaxis.set_label_position("right")
ax_lr.plot(df["index"], df["lr"], label="LR", linestyle='--', alpha=0.3, color="tab:brown")
ax_acc.plot(df["index"], df["accuracy"], label="Accuracy", linestyle='--', alpha=0.3, color="tab:green")
ax_acc.plot(df["index"], df["avg_accuracy"], label="Avg Accuracy", color="tab:green")
ax_loss.plot(df["index"], df["loss"], label="Loss", linestyle=':', alpha=0.3, color="tab:blue")
ax_loss.plot(df["index"], df["ema_loss"], label="EMA Loss", linestyle='--', alpha=0.7, color="tab:blue")
ax_loss.plot(df["index"], df["avg_loss"], label=f"Avg Loss", color="tab:blue")
ax_time.plot(df["index"], df["queue_ms"], label="Queueing", linestyle=':', alpha=0.3, color="#ffa64d")
ax_time.plot(df["index"], df["compute_ms"], label="Computing", linestyle=':', alpha=0.3, color="#cc6600")
ax_time.plot(df["index"], df["infer_ms"], label="Inferring", linestyle=':', alpha=0.3, color="#ffb84d")
ax_time.plot(df["index"], df["total_ms"], label="Total Time",linestyle='--', alpha=0.7, color="tab:orange")
ax_time.plot(df["index"], df["avg_total_ms"], label="Avg Total Time", color="tab:orange")
# Add epoch/batch label ticks to secondary X axis
tick_locs = ax_time.get_xticks()
tick_labels = []
if not df.empty:
for loc in tick_locs:
idx = (abs(df["index"] - loc)).argmin()
e, b = int(df.iloc[idx]["epoch"]), int(df.iloc[idx]["batch"])
tick_labels.append(f"{e + 1}/{b + 1}")
else:
tick_labels = [""] * len(tick_locs)
ax_loss.set_xlim(ax_time.get_xlim())
ax_loss.set_xticks(tick_locs)
ax_loss.set_xticklabels(tick_labels)
# Gather all legend entries
lines = []
labels = []
for ax in [ax_lr, ax_acc, ax_loss]:
l, lab = ax.get_legend_handles_labels()
lines += l
labels += lab
# Add combined legend to ax_loss
ax_loss.legend(lines, labels, loc='upper left') # or 'upper right', etc.
ax_time.legend(loc='upper left')
ani = animation.FuncAnimation(fig, animate, interval=1000)
plt.show()