File size: 3,683 Bytes
fbd9366 | 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 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 | #!/usr/bin/env python3
"""Compare loss curves between LoRA and full fine-tuning runs.
Usage:
python scripts/compare_loss.py \
--lora-log ./checkpoints/dreamzero_droid_lora/loss_log.jsonl \
--full-log ./checkpoints/dreamzero_droid_full_finetune/loss_log.jsonl \
[--plot loss_comparison.png]
"""
import argparse
import json
def load_loss_log(path):
entries = []
with open(path) as f:
for line in f:
line = line.strip()
if line:
entries.append(json.loads(line))
return entries
def print_comparison_table(lora_entries, full_entries):
# Index by step
lora_by_step = {e["step"]: e for e in lora_entries}
full_by_step = {e["step"]: e for e in full_entries}
all_steps = sorted(set(lora_by_step.keys()) | set(full_by_step.keys()))
header = f"{'Step':>6} {'LoRA Loss':>10} {'Full Loss':>10} {'LoRA Dyn':>10} {'Full Dyn':>10} {'LoRA Act':>10} {'Full Act':>10}"
print(header)
print("-" * len(header))
for step in all_steps:
lora = lora_by_step.get(step, {})
full = full_by_step.get(step, {})
def fmt(d, key):
v = d.get(key)
return f"{v:10.4f}" if v is not None else f"{'—':>10}"
print(
f"{step:>6} "
f"{fmt(lora, 'loss')} {fmt(full, 'loss')} "
f"{fmt(lora, 'dynamics_loss_avg')} {fmt(full, 'dynamics_loss_avg')} "
f"{fmt(lora, 'action_loss_avg')} {fmt(full, 'action_loss_avg')}"
)
def plot_comparison(lora_entries, full_entries, output_path):
try:
import matplotlib.pyplot as plt
except ImportError:
print("matplotlib not installed, skipping plot generation.")
print("Install with: pip install matplotlib")
return
metrics = [
("loss", "Total Loss"),
("dynamics_loss_avg", "Dynamics Loss"),
("action_loss_avg", "Action Loss"),
]
fig, axes = plt.subplots(1, len(metrics), figsize=(5 * len(metrics), 4))
if len(metrics) == 1:
axes = [axes]
for ax, (key, title) in zip(axes, metrics):
lora_steps = [e["step"] for e in lora_entries if key in e]
lora_vals = [e[key] for e in lora_entries if key in e]
full_steps = [e["step"] for e in full_entries if key in e]
full_vals = [e[key] for e in full_entries if key in e]
if lora_steps:
ax.plot(lora_steps, lora_vals, label="LoRA", marker="o", markersize=3)
if full_steps:
ax.plot(full_steps, full_vals, label="Full FT", marker="s", markersize=3)
ax.set_title(title)
ax.set_xlabel("Step")
ax.set_ylabel("Loss")
ax.legend()
ax.grid(True, alpha=0.3)
fig.tight_layout()
fig.savefig(output_path, dpi=150)
print(f"Plot saved to {output_path}")
def main():
parser = argparse.ArgumentParser(description="Compare LoRA vs full fine-tuning loss curves")
parser.add_argument("--lora-log", required=True, help="Path to LoRA run loss_log.jsonl")
parser.add_argument("--full-log", required=True, help="Path to full FT run loss_log.jsonl")
parser.add_argument("--plot", default=None, help="Output path for comparison plot (e.g., loss_comparison.png)")
args = parser.parse_args()
lora_entries = load_loss_log(args.lora_log)
full_entries = load_loss_log(args.full_log)
print(f"LoRA: {len(lora_entries)} log entries")
print(f"Full: {len(full_entries)} log entries")
print()
print_comparison_table(lora_entries, full_entries)
if args.plot:
plot_comparison(lora_entries, full_entries, args.plot)
if __name__ == "__main__":
main()
|