| """
|
| Gera graficos de treinamento do ResearcherBot-1 Multimodal.
|
|
|
| Produz 4 graficos:
|
| 1. Training & Validation Loss
|
| 2. Perplexity (Train & Val)
|
| 3. Learning Rate Schedule
|
| 4. MoE Auxiliary Loss (Validation)
|
|
|
| Salva em: training_plots/
|
| """
|
|
|
| import json
|
| import os
|
| import matplotlib
|
| matplotlib.use("Agg")
|
| import matplotlib.pyplot as plt
|
| import matplotlib.ticker as mticker
|
|
|
|
|
| plt.rcParams.update({
|
| "figure.facecolor": "#0d1117",
|
| "axes.facecolor": "#161b22",
|
| "axes.edgecolor": "#30363d",
|
| "axes.labelcolor": "#c9d1d9",
|
| "xtick.color": "#8b949e",
|
| "ytick.color": "#8b949e",
|
| "text.color": "#c9d1d9",
|
| "grid.color": "#21262d",
|
| "grid.linestyle": "--",
|
| "grid.alpha": 0.6,
|
| "font.family": "sans-serif",
|
| "font.size": 11,
|
| "legend.facecolor": "#161b22",
|
| "legend.edgecolor": "#30363d",
|
| "legend.fontsize": 10,
|
| })
|
|
|
| COLORS = {
|
| "train": "#58a6ff",
|
| "val": "#f78166",
|
| "lr": "#7ee787",
|
| "aux": "#d2a8ff",
|
| "accent": "#ffa657",
|
| }
|
|
|
| OUTPUT_DIR = "training_plots"
|
|
|
|
|
| def load_metrics(path="training_metrics.json"):
|
| with open(path, "r", encoding="utf-8") as f:
|
| return json.load(f)
|
|
|
|
|
| def plot_loss(m, save_dir):
|
| """Grafico 1: Train & Val Loss."""
|
| fig, ax = plt.subplots(figsize=(10, 5.5))
|
|
|
| ax.plot(m["train_steps"], m["train_loss"], color=COLORS["train"],
|
| linewidth=1.8, label="Train Loss", alpha=0.9)
|
| ax.plot(m["val_steps"], m["val_loss"], color=COLORS["val"],
|
| linewidth=2.2, marker="o", markersize=7, label="Val Loss",
|
| linestyle="--", alpha=0.9)
|
|
|
| ax.set_xlabel("Step")
|
| ax.set_ylabel("Loss")
|
| ax.set_title("Training & Validation Loss", fontsize=14, fontweight="bold", pad=12)
|
| ax.legend(loc="upper right")
|
| ax.grid(True)
|
| ax.set_ylim(bottom=0)
|
|
|
|
|
| best_idx = m["val_loss"].index(min(m["val_loss"]))
|
| best_step = m["val_steps"][best_idx]
|
| best_val = m["val_loss"][best_idx]
|
| ax.annotate(f"Best: {best_val:.4f}",
|
| xy=(best_step, best_val),
|
| xytext=(best_step + 100, best_val + 0.5),
|
| arrowprops=dict(arrowstyle="->", color=COLORS["accent"], lw=1.5),
|
| fontsize=10, color=COLORS["accent"], fontweight="bold")
|
|
|
| plt.tight_layout()
|
| path = os.path.join(save_dir, "loss.png")
|
| fig.savefig(path, dpi=150, bbox_inches="tight")
|
| plt.close(fig)
|
| print(f" Salvo: {path}")
|
|
|
|
|
| def plot_perplexity(m, save_dir):
|
| """Grafico 2: Train & Val Perplexity."""
|
| fig, ax = plt.subplots(figsize=(10, 5.5))
|
|
|
| ax.plot(m["train_steps"], m["train_ppl"], color=COLORS["train"],
|
| linewidth=1.8, label="Train PPL", alpha=0.9)
|
| ax.plot(m["val_steps"], m["val_ppl"], color=COLORS["val"],
|
| linewidth=2.2, marker="s", markersize=7, label="Val PPL",
|
| linestyle="--", alpha=0.9)
|
|
|
| ax.set_xlabel("Step")
|
| ax.set_ylabel("Perplexity")
|
| ax.set_title("Perplexity (Train & Validation)", fontsize=14, fontweight="bold", pad=12)
|
| ax.legend(loc="upper right")
|
| ax.grid(True)
|
| ax.set_yscale("log")
|
|
|
|
|
| final_train_ppl = m["train_ppl"][-1]
|
| ax.axhline(y=final_train_ppl, color=COLORS["train"], linestyle=":", alpha=0.4)
|
| ax.text(m["train_steps"][-1], final_train_ppl * 1.1,
|
| f"Final: {final_train_ppl:.1f}",
|
| color=COLORS["train"], fontsize=9, ha="right")
|
|
|
| plt.tight_layout()
|
| path = os.path.join(save_dir, "perplexity.png")
|
| fig.savefig(path, dpi=150, bbox_inches="tight")
|
| plt.close(fig)
|
| print(f" Salvo: {path}")
|
|
|
|
|
| def plot_learning_rate(m, save_dir):
|
| """Grafico 3: Learning Rate Schedule."""
|
| fig, ax = plt.subplots(figsize=(10, 5.5))
|
|
|
| ax.plot(m["train_steps"], m["lr"], color=COLORS["lr"],
|
| linewidth=2.0, alpha=0.9)
|
| ax.fill_between(m["train_steps"], 0, m["lr"],
|
| color=COLORS["lr"], alpha=0.15)
|
|
|
| ax.set_xlabel("Step")
|
| ax.set_ylabel("Learning Rate")
|
| ax.set_title("Learning Rate Schedule (Warmup + Cosine Decay)",
|
| fontsize=14, fontweight="bold", pad=12)
|
| ax.grid(True)
|
| ax.yaxis.set_major_formatter(mticker.FormatStrFormatter("%.1e"))
|
|
|
|
|
| warmup_end = 500
|
| ax.axvline(x=warmup_end, color=COLORS["accent"], linestyle="--",
|
| alpha=0.6, linewidth=1.5)
|
| ax.text(warmup_end + 20, max(m["lr"]) * 0.95, "Warmup\nEnd",
|
| color=COLORS["accent"], fontsize=9, fontweight="bold")
|
|
|
| plt.tight_layout()
|
| path = os.path.join(save_dir, "learning_rate.png")
|
| fig.savefig(path, dpi=150, bbox_inches="tight")
|
| plt.close(fig)
|
| print(f" Salvo: {path}")
|
|
|
|
|
| def plot_aux_loss(m, save_dir):
|
| """Grafico 4: MoE Auxiliary Loss."""
|
| if not m.get("val_aux_loss"):
|
| return
|
|
|
| fig, ax = plt.subplots(figsize=(10, 5.5))
|
|
|
| ax.plot(m["val_steps"], m["val_aux_loss"], color=COLORS["aux"],
|
| linewidth=2.2, marker="D", markersize=7, alpha=0.9)
|
| ax.fill_between(m["val_steps"], 0, m["val_aux_loss"],
|
| color=COLORS["aux"], alpha=0.1)
|
|
|
| ax.set_xlabel("Step")
|
| ax.set_ylabel("Aux Loss (Load Balance)")
|
| ax.set_title("MoE Load Balancing Loss (Validation)",
|
| fontsize=14, fontweight="bold", pad=12)
|
| ax.grid(True)
|
|
|
|
|
| ax.axhline(y=1.0, color="#8b949e", linestyle=":", alpha=0.5, linewidth=1.2)
|
| ax.text(m["val_steps"][-1], 1.05, "Ideal (uniforme)",
|
| color="#8b949e", fontsize=9, ha="right")
|
|
|
| plt.tight_layout()
|
| path = os.path.join(save_dir, "moe_aux_loss.png")
|
| fig.savefig(path, dpi=150, bbox_inches="tight")
|
| plt.close(fig)
|
| print(f" Salvo: {path}")
|
|
|
|
|
| def plot_combined_dashboard(m, save_dir):
|
| """Grafico 5: Dashboard combinado 2x2."""
|
| fig, axes = plt.subplots(2, 2, figsize=(14, 10))
|
|
|
|
|
| ax = axes[0, 0]
|
| ax.plot(m["train_steps"], m["train_loss"], color=COLORS["train"],
|
| linewidth=1.5, label="Train", alpha=0.9)
|
| ax.plot(m["val_steps"], m["val_loss"], color=COLORS["val"],
|
| linewidth=2, marker="o", markersize=5, label="Val",
|
| linestyle="--", alpha=0.9)
|
| ax.set_xlabel("Step")
|
| ax.set_ylabel("Loss")
|
| ax.set_title("Loss", fontsize=12, fontweight="bold")
|
| ax.legend(fontsize=9)
|
| ax.grid(True)
|
| ax.set_ylim(bottom=0)
|
|
|
|
|
| ax = axes[0, 1]
|
| ax.plot(m["train_steps"], m["train_ppl"], color=COLORS["train"],
|
| linewidth=1.5, label="Train", alpha=0.9)
|
| ax.plot(m["val_steps"], m["val_ppl"], color=COLORS["val"],
|
| linewidth=2, marker="s", markersize=5, label="Val",
|
| linestyle="--", alpha=0.9)
|
| ax.set_xlabel("Step")
|
| ax.set_ylabel("Perplexity")
|
| ax.set_title("Perplexity", fontsize=12, fontweight="bold")
|
| ax.legend(fontsize=9)
|
| ax.grid(True)
|
| ax.set_yscale("log")
|
|
|
|
|
| ax = axes[1, 0]
|
| ax.plot(m["train_steps"], m["lr"], color=COLORS["lr"], linewidth=1.8, alpha=0.9)
|
| ax.fill_between(m["train_steps"], 0, m["lr"], color=COLORS["lr"], alpha=0.1)
|
| ax.set_xlabel("Step")
|
| ax.set_ylabel("LR")
|
| ax.set_title("Learning Rate", fontsize=12, fontweight="bold")
|
| ax.grid(True)
|
| ax.yaxis.set_major_formatter(mticker.FormatStrFormatter("%.1e"))
|
|
|
|
|
| ax = axes[1, 1]
|
| if m.get("val_aux_loss"):
|
| ax.plot(m["val_steps"], m["val_aux_loss"], color=COLORS["aux"],
|
| linewidth=2, marker="D", markersize=5, alpha=0.9)
|
| ax.axhline(y=1.0, color="#8b949e", linestyle=":", alpha=0.4)
|
| ax.set_xlabel("Step")
|
| ax.set_ylabel("Aux Loss")
|
| ax.set_title("MoE Aux Loss", fontsize=12, fontweight="bold")
|
| ax.grid(True)
|
|
|
| fig.suptitle("ResearcherBot-1 Multimodal | 19.9M params | MoE 8 experts top-2",
|
| fontsize=15, fontweight="bold", y=0.98, color="#58a6ff")
|
| plt.tight_layout(rect=[0, 0, 1, 0.96])
|
| path = os.path.join(save_dir, "dashboard.png")
|
| fig.savefig(path, dpi=150, bbox_inches="tight")
|
| plt.close(fig)
|
| print(f" Salvo: {path}")
|
|
|
|
|
| def main():
|
| os.makedirs(OUTPUT_DIR, exist_ok=True)
|
|
|
| print("=" * 50)
|
| print(" Gerando graficos de treinamento...")
|
| print("=" * 50)
|
|
|
| m = load_metrics()
|
| print(f" Steps: {m['total_steps']} | Tempo: {m['total_time_min']:.1f} min")
|
| print(f" Best Val Loss: {m['best_val_loss']:.4f}")
|
| print()
|
|
|
| plot_loss(m, OUTPUT_DIR)
|
| plot_perplexity(m, OUTPUT_DIR)
|
| plot_learning_rate(m, OUTPUT_DIR)
|
| plot_aux_loss(m, OUTPUT_DIR)
|
| plot_combined_dashboard(m, OUTPUT_DIR)
|
|
|
| print(f"\n Todos os graficos em: {OUTPUT_DIR}/")
|
|
|
|
|
| if __name__ == "__main__":
|
| main()
|
|
|
