plot_training_convergence

Plot training convergence: up to three subplots (mean_by_class, mean_by_feature_class, correlation).

Data source: exactly one of trainer, model_path, or training_stats.

• trainer: uses trainer.get_training_stats() (or in-memory stats if available).
• model_path: uses load_training_stats(model_path).
• training_stats: dict with keys training_stats, best_score_checkpoint (or raw training_stats dict).

Three plot options, each in its own subplot when enabled:

• plot_mean_by_class: mean encoded value per label over steps.
• plot_mean_by_feature_class: mean encoded value per feature class over steps.
• plot_correlation: correlation over steps. Best checkpoint step is marked in each subplot.
• class_spread: shade encoded value spread per class behind each mean line

("minmax" = min-max, "iqr" = Q1-Q3, "ci95" = 95% confidence interval).

Parameters:

Name	Type	Description	Default
trainer	Any	Trainer instance with get_training_stats(model_path) or similar.	None
model_path	Optional[str]	Path to saved model dir (training.json).	None
training_stats	Optional[Dict[str, Any]]	Pre-loaded run info or raw training_stats dict.	None
plot_mean_by_class	bool	Add a subplot for mean_by_class.	True
plot_mean_by_feature_class	Optional[bool]	Add a subplot for mean_by_feature_class. None = auto (False when redundant with mean_by_class, True otherwise).	None
plot_correlation	bool	Add a subplot for correlation.	True
class_spread	ClassSpreadMode	Shade spread per class (and feature class) behind mean lines. "minmax": min-max band. "iqr": interquartile range (Q1-Q3). "ci95": 95% confidence interval around the mean (mean +/- 1.96 SE). None disables spread shading. Requires matching keys in training stats (recorded from new runs).	None
best_step	bool	Mark the best checkpoint step (vertical line + point on correlation).	True
label_name_mapping	Optional[Dict[str, str]]	Optional display names for label values.	None
output	Optional[str]	Explicit output path for the plot file.	None
experiment_dir	Optional[str]	Used with resolve_output_path for default artifact path.	None
show	bool	Whether to call plt.show().	True
title	Union[str, bool]	True (default run_id), False, or custom string.	True
figsize	Optional[Tuple[float, float]]	(width, height) for figure.	None
img_format	str	File extension used when resolving the default output path.	'png'
dpi	Optional[int]	Optional Matplotlib savefig DPI.	None
legend_ncol	Optional[int]	Number of columns for the external legend when there are >= 6 series (default 1).	None
legend_bbox_to_anchor	Optional[Tuple[float, float]]	(x, y) for the external legend when >= 6 series (default (1.02, 0.5)).	None
legend_loc	Optional[str]	Matplotlib loc for the external legend when >= 6 series (default "center left").	None
highlight_non_convergence	Optional[bool]	When True, append a non-convergence marker to the title for non-converged runs. None uses TrainingArguments.highlight_non_convergence.	None
return_fig_ax	bool	If True, return (fig, axes) and leave the figure open so callers can customize it before showing, saving again, or closing it. Existing saving/display behavior still runs when output/experiment_dir or show are set.	False
**kwargs	Any	Reserved for compatibility with trainer visualizer wrappers.	{}

Returns:

Type	Description
Any	Path to saved plot file, or "" if nothing to plot or no path. If return_fig_ax=True,
Any	returns (fig, axes).

Source code in gradiend/visualizer/convergence.py

def plot_training_convergence(
    trainer: Any = None,
    model_path: Optional[str] = None,
    training_stats: Optional[Dict[str, Any]] = None,
    *,
    plot_mean_by_class: bool = True,
    plot_mean_by_feature_class: Optional[bool] = None,
    plot_correlation: bool = True,
    class_spread: ClassSpreadMode = None,
    best_step: bool = True,
    label_name_mapping: Optional[Dict[str, str]] = None,
    output: Optional[str] = None,
    experiment_dir: Optional[str] = None,
    show: bool = True,
    title: Union[str, bool] = True,
    figsize: Optional[Tuple[float, float]] = None,
    img_format: str = "png",
    dpi: Optional[int] = None,
    legend_ncol: Optional[int] = None,
    legend_bbox_to_anchor: Optional[Tuple[float, float]] = None,
    legend_loc: Optional[str] = None,
    highlight_non_convergence: Optional[bool] = None,
    return_fig_ax: bool = False,
    **kwargs: Any,
) -> Any:
    """
    Plot training convergence: up to three subplots (mean_by_class, mean_by_feature_class, correlation).

    Data source: exactly one of trainer, model_path, or training_stats.

    - trainer: uses trainer.get_training_stats() (or in-memory stats if available).
    - model_path: uses load_training_stats(model_path).
    - training_stats: dict with keys training_stats, best_score_checkpoint (or raw training_stats dict).

    Three plot options, each in its own subplot when enabled:

    - plot_mean_by_class: mean encoded value per label over steps.
    - plot_mean_by_feature_class: mean encoded value per feature class over steps.
    - plot_correlation: correlation over steps. Best checkpoint step is marked in each subplot.
    - class_spread: shade encoded value spread per class behind each mean line

      (``"minmax"`` = min-max, ``"iqr"`` = Q1-Q3, ``"ci95"`` = 95% confidence interval).

    Args:
        trainer: Trainer instance with get_training_stats(model_path) or similar.
        model_path: Path to saved model dir (training.json).
        training_stats: Pre-loaded run info or raw training_stats dict.
        plot_mean_by_class: Add a subplot for mean_by_class.
        plot_mean_by_feature_class: Add a subplot for mean_by_feature_class.
            None = auto (False when redundant with mean_by_class, True otherwise).
        plot_correlation: Add a subplot for correlation.
        class_spread: Shade spread per class (and feature class) behind mean lines.
            ``"minmax"``: min-max band. ``"iqr"``: interquartile range (Q1-Q3).
            ``"ci95"``: 95% confidence interval around the mean (mean +/- 1.96 SE).
            ``None`` disables spread shading.
            Requires matching keys in training stats (recorded from new runs).
        best_step: Mark the best checkpoint step (vertical line + point on correlation).
        label_name_mapping: Optional display names for label values.
        output: Explicit output path for the plot file.
        experiment_dir: Used with resolve_output_path for default artifact path.
        show: Whether to call plt.show().
        title: True (default run_id), False, or custom string.
        figsize: (width, height) for figure.
        img_format: File extension used when resolving the default output path.
        dpi: Optional Matplotlib savefig DPI.
        legend_ncol: Number of columns for the external legend when there are >= 6 series (default 1).
        legend_bbox_to_anchor: (x, y) for the external legend when >= 6 series (default (1.02, 0.5)).
        legend_loc: Matplotlib loc for the external legend when >= 6 series (default "center left").
        highlight_non_convergence: When True, append a non-convergence marker to the title for
            non-converged runs. ``None`` uses ``TrainingArguments.highlight_non_convergence``.
        return_fig_ax: If True, return ``(fig, axes)`` and leave the figure open so callers can
            customize it before showing, saving again, or closing it. Existing saving/display
            behavior still runs when ``output``/``experiment_dir`` or ``show`` are set.
        **kwargs: Reserved for compatibility with trainer visualizer wrappers.

    Returns:
        Path to saved plot file, or "" if nothing to plot or no path. If ``return_fig_ax=True``,
        returns ``(fig, axes)``.
    """
    if training_stats is not None:
        if "training_stats" in training_stats and "best_score_checkpoint" in training_stats:
            run_info = training_stats
        else:
            run_info = {"training_stats": training_stats, "best_score_checkpoint": {}}
    elif model_path:
        from gradiend.trainer.core.stats import load_training_stats

        run_info = load_training_stats(model_path)
        if run_info is None:
            logger.warning("No training.json at %s", model_path)
            return ""
    elif trainer is not None:
        get_stats = getattr(trainer, "get_training_stats", None)
        if get_stats is not None:
            run_info = get_stats()
        else:
            run_info = None
        if run_info is None:
            logger.warning("Could not get training stats from trainer")
            return ""
    else:
        raise ValueError("Provide one of trainer, model_path, or training_stats")

    plt = _require_matplotlib()

    ts = run_info.get("training_stats") or run_info
    if isinstance(ts, dict) and "training_stats" in ts:
        ts = ts["training_stats"]
    if not ts:
        logger.warning("No training_stats to plot")
        return ""

    bsc = run_info.get("best_score_checkpoint") or {}
    best_step_val = bsc.get("global_step")
    if best_step_val is not None:
        try:
            best_step_val = int(best_step_val)
        except (TypeError, ValueError):
            best_step_val = None

    steps, series_by_class, series_by_fc = _steps_and_values(
        ts, mean_by_class=plot_mean_by_class, mean_by_feature_class=plot_mean_by_feature_class is not False
    )
    if plot_mean_by_feature_class is None:
        plot_mean_by_feature_class = not _is_mean_by_feature_class_redundant(ts, series_by_class, series_by_fc)
    corr_series = _correlation_series(ts) if plot_correlation else []

    has_mbc = plot_mean_by_class and bool(series_by_class)
    has_mbfc = plot_mean_by_feature_class and bool(series_by_fc)
    has_corr = plot_correlation and bool(corr_series)
    if not (has_mbc or has_mbfc or has_corr):
        logger.warning("No plottable series in training_stats")
        return ""

    n_sub = (1 if has_mbc else 0) + (1 if has_mbfc else 0) + (1 if has_corr else 0)
    n_legend_entries = max(len(series_by_class) if series_by_class else 0, len(series_by_fc) if series_by_fc else 0)
    # When many legend entries (e.g. identity transitions), use larger height so subplots are not squashed by legend
    height_per_sub = 2.0 if n_legend_entries >= 6 else 1.5
    fig, axes = plt.subplots(n_sub, 1, sharex=True, figsize=figsize or (6, height_per_sub * n_sub))
    if n_sub == 1:
        axes = [axes]
    # Leave space for figure-level legend when >= 6 series
    if n_legend_entries >= 6:
        fig.subplots_adjust(right=0.72)

    best_corr_val = bsc.get("correlation")
    if best_corr_val is not None:
        try:
            best_corr_val = float(best_corr_val)
        except (TypeError, ValueError):
            best_corr_val = None

    draw_convergence_axes(
        ts,
        axes,
        best_step_val=best_step_val if best_step else None,
        best_corr=best_corr_val if best_step else None,
        label_name_mapping=label_name_mapping,
        plot_mean_by_class=has_mbc,
        plot_mean_by_feature_class=has_mbfc,
        plot_correlation=has_corr,
        class_spread=class_spread,
        legend_ncol=legend_ncol,
        legend_bbox_to_anchor=legend_bbox_to_anchor,
        legend_loc=legend_loc,
    )

    highlight = resolve_highlight_non_convergence(highlight_non_convergence, trainer=trainer)
    resolved_title = resolve_plot_title_with_convergence(
        title,
        trainer=trainer,
        run_info=run_info,
        highlight_non_convergence=highlight,
    )
    if resolved_title is not False:
        fig.suptitle(str(resolved_title), fontsize=10)
    plt.tight_layout()

    out_path = None
    if output:
        out_path = output
    else:
        exp_dir = experiment_dir or (getattr(trainer, "experiment_dir", None) if trainer is not None else None)
        out_path = resolve_output_path(exp_dir, None, ARTIFACT_CONVERGENCE_PLOT)
    if out_path and img_format:
        ext = img_format if img_format.startswith(".") else f".{img_format}"
        out_path = os.path.splitext(out_path)[0] + ext

    if out_path is None and not show and not return_fig_ax:
        raise ValueError(
            "output is required when experiment_dir is not set and not show=True."
        )

    if out_path:
        os.makedirs(os.path.dirname(out_path) or ".", exist_ok=True)
        save_kwargs: Dict[str, Any] = {"bbox_inches": "tight"}
        if dpi is not None:
            save_kwargs["dpi"] = dpi
        plt.savefig(out_path, **save_kwargs)
        logger.info("Saved convergence plot: %s", out_path)
    if show:
        plt.show()
    if return_fig_ax:
        return fig, axes
    plt.close(fig)
    return out_path or ""