| |
| """Plot training loss curves from training logs for comparison.""" |
|
|
| import re |
| import sys |
| import matplotlib |
| matplotlib.use('Agg') |
| import matplotlib.pyplot as plt |
| import numpy as np |
|
|
| def parse_loss_from_log(log_path): |
| """Extract (step, loss, grad_norm) from training log.""" |
| steps, losses, grad_norms = [], [], [] |
| pattern = re.compile(r"'loss': ([\d.]+), 'grad_norm': ([\d.]+)") |
| with open(log_path) as f: |
| for line in f: |
| m = pattern.search(line) |
| if m: |
| steps.append(len(losses) + 1) |
| losses.append(float(m.group(1))) |
| grad_norms.append(float(m.group(2))) |
| return steps, losses, grad_norms |
|
|
|
|
| def parse_loss_from_tensorboard(tb_dir): |
| """Extract (step, loss, grad_norm) from tensorboard event files.""" |
| try: |
| from tensorboard.backend.event_processing.event_accumulator import EventAccumulator |
| except ImportError: |
| print("tensorboard not installed, skipping TB parsing") |
| return [], [], [] |
|
|
| ea = EventAccumulator(tb_dir) |
| ea.Reload() |
|
|
| steps, losses, grad_norms = [], [], [] |
| if 'train/loss' in ea.scalars.Keys(): |
| for event in ea.scalars.Items('train/loss'): |
| steps.append(event.step) |
| losses.append(event.value) |
| elif 'loss' in ea.scalars.Keys(): |
| for event in ea.scalars.Items('loss'): |
| steps.append(event.step) |
| losses.append(event.value) |
|
|
| if 'train/grad_norm' in ea.scalars.Keys(): |
| for event in ea.scalars.Items('train/grad_norm'): |
| grad_norms.append(event.value) |
| elif 'grad_norm' in ea.scalars.Keys(): |
| for event in ea.scalars.Items('grad_norm'): |
| grad_norms.append(event.value) |
|
|
| return steps, losses, grad_norms |
|
|
|
|
| def plot_comparison(curves, output_path, title="Training Loss Comparison"): |
| """Plot loss and grad_norm for multiple runs.""" |
| fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5)) |
|
|
| for name, (steps, losses, grad_norms) in curves.items(): |
| if losses: |
| ax1.plot(steps[:len(losses)], losses, label=name, alpha=0.8) |
| if grad_norms: |
| ax2.plot(steps[:len(grad_norms)], grad_norms, label=name, alpha=0.8) |
|
|
| ax1.set_xlabel('Step') |
| ax1.set_ylabel('Loss') |
| ax1.set_title('Training Loss') |
| ax1.legend() |
| ax1.grid(True, alpha=0.3) |
|
|
| ax2.set_xlabel('Step') |
| ax2.set_ylabel('Grad Norm') |
| ax2.set_title('Gradient Norm') |
| ax2.legend() |
| ax2.grid(True, alpha=0.3) |
|
|
| plt.suptitle(title, fontsize=14, fontweight='bold') |
| plt.tight_layout() |
| plt.savefig(output_path, dpi=150, bbox_inches='tight') |
| print(f"Saved: {output_path}") |
|
|
|
|
| if __name__ == '__main__': |
| import os |
| import glob |
|
|
| curves = {} |
|
|
| |
| baseline_tb_dirs = sorted(glob.glob('/mnt/bn/leonworkspace/terry/model/qwen3vl-4b-roi-K24T3-185k-ddp-verify/runs/*')) |
| if baseline_tb_dirs: |
| steps, losses, grad_norms = parse_loss_from_tensorboard(baseline_tb_dirs[-1]) |
| if losses: |
| curves['baseline (ddp-verify)'] = (steps, losses, grad_norms) |
|
|
| |
| bidir_log = '/tmp/bidir_train.log' |
| if os.path.exists(bidir_log): |
| steps, losses, grad_norms = parse_loss_from_log(bidir_log) |
| if losses: |
| curves['bidir (all_visual)'] = (steps, losses, grad_norms) |
|
|
| |
| bidir_tb_dirs = sorted(glob.glob('/mnt/bn/leonworkspace/terry/model/qwen3vl-4b-roi-K24T3-185k-bidir/runs/*')) |
| if bidir_tb_dirs: |
| steps, losses, grad_norms = parse_loss_from_tensorboard(bidir_tb_dirs[-1]) |
| if losses and len(losses) > 10: |
| curves['bidir (tensorboard)'] = (steps, losses, grad_norms) |
|
|
| if not curves: |
| print("No data found!") |
| sys.exit(1) |
|
|
| output_path = '/opt/tiger/thothvl_pretrain/visualize/bidir_vs_baseline_loss.png' |
| plot_comparison(curves, output_path, "Bidir Visual Attention Training: Loss Comparison") |
|
|
