metrics_explorer.py

# run with:
# streamlit run metrics_explorer.py
import json
from pathlib import Path
from typing import Any, Dict, List, Tuple, Union

import altair as alt
import numpy as np
import pandas as pd
import streamlit as st

st.set_page_config(
    page_title="FSL metrics",
    page_icon=None,
    layout="wide",
)

"""# FSLH - metrics explorer"""

root = Path(
    "/home/mprzewie/coding/gmum_mnt/uj/few-shot-hypernets/save/checkpoints/cross_char/"
)

METRICS_FILE = "metrics.json"
ARGS_FILE = "args.json"

LOADING_PROGRESS = st.progress(0.0)

PROGRESS_TEXT = st.empty()
HN_PREFIX = "hn"


def metrics_dict_to_df(
    experiment_name: str,
    metrics_dict: Dict[str, List[Union[float, List[float]]]],
    args_dict: Dict[str, Any],
) -> pd.DataFrame:
    rows = []
    for m_name, values in metrics_dict.items():
        for e, vls in enumerate(values):
            vls = [vls] if not isinstance(vls, list) else [np.mean(vls)]  # vls
            rows.extend(
                [
                    {
                        "exp_name": experiment_name,
                        "met_name": m_name,
                        "epoch": e,
                        "value": v,
                        **{
                            ak: av
                            for (ak, av) in args_dict.items()
                            if ak.startswith(HN_PREFIX)
                        },
                    }
                    for v in vls
                ]
            )
    return pd.DataFrame(rows)


@st.cache(allow_output_mutation=True, suppress_st_warning=True)
def load_experiment(experiment_path: Path) -> Tuple[Dict, pd.DataFrame]:
    with (experiment_path / METRICS_FILE).open("r") as f:
        metrics = json.load(f)
    with (experiment_path / ARGS_FILE).open("r") as f:
        args = json.load(f)
    return args, metrics_dict_to_df(
        experiment_name=experiment_path.name, metrics_dict=metrics, args_dict=args
    )


loggable_experiments = {
    p.parent.name: p
    for p in root.glob(f"*/{METRICS_FILE}")
    if (p.parent / ARGS_FILE).exists()
}

experiment_args = dict()
experiment_metrics = dict()

loggable_experiments = list(loggable_experiments.items())  # [:3]

for i, (e, p) in enumerate(loggable_experiments):
    PROGRESS_TEXT.text(f"loading {p.parent.name} {i}/{len(loggable_experiments)}")

    try:
        args, metrics_df = load_experiment(p.parent)
        experiment_args[e] = args
        experiment_metrics[e] = metrics_df

        LOADING_PROGRESS.progress((i + 1) / len(loggable_experiments))
    except Exception as exc:
        print(exc)
        pass

PROGRESS_TEXT.text(f"Loaded {len(experiment_metrics)} experiments")

df = pd.concat([mdf for mdf in experiment_metrics.values()])

available_metrics = sorted(df.met_name.unique())
all_args = sorted({a for ad in experiment_args.values() for a in ad.keys()})


"""## Selected metric over the course of epochs"""

selected_metric = st.selectbox(
    "Select metric",
    available_metrics,
    index=available_metrics.index("accuracy_val_max"),
)
aggregate_y = st.checkbox("Aggregate Y?", value=True)

st.altair_chart(
    alt.Chart(
        df[df.met_name == selected_metric],
    )
    .mark_line(point=True)
    .encode(
        x="epoch",
        y=alt.Y("value", aggregate=("mean" if aggregate_y else alt.Undefined)),
        color="exp_name",
        tooltip=["exp_name", "value", "epoch"]
        + [a for a in all_args if a.startswith(HN_PREFIX)],
    )
    .configure_legend(labelLimit=0)
    .interactive()
    .properties(title=selected_metric),
    use_container_width=True,
)


"""## How do hyperparams influence the metric?"""

for a in all_args:
    if a.startswith(HN_PREFIX):
        unique_as = sorted(df[a].unique())
        with st.expander(f"{a} := {unique_as}", expanded=(len(unique_as) > 1)):
            st.altair_chart(
                alt.Chart(
                    df[df.met_name == selected_metric],
                )
                .mark_line(point=True)
                .encode(
                    x="epoch",
                    y=alt.Y(
                        "value", aggregate=("mean" if aggregate_y else alt.Undefined)
                    ),
                    color="exp_name",
                    tooltip=["exp_name", "value", "epoch"]
                    + [a for a in all_args if a.startswith(HN_PREFIX)],
                    column=a,
                )
                .configure_legend(labelLimit=0)
                .interactive()
                .properties(title=selected_metric)
            )