Add val/train loss to evaluation and plots

- _evaluate() now computes per-batch loss when loss_fn is provided
- Val loss and train recall loss saved in val.csv and train_recall.csv
- Overview plot shows train vs val loss curves side by side
- Helps detect overfitting: val loss diverging from train loss

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
pikaliov
2026-04-22 08:05:12 +03:00
parent 93ad66810d
commit df60e83ead
2 changed files with 31 additions and 6 deletions

View File

@@ -114,12 +114,15 @@ def plot_combined(train_df: pd.DataFrame, val_df: pd.DataFrame, out_dir: Path) -
fig, axes = plt.subplots(1, 3, figsize=(18, 5))
fig.suptitle("Training Overview", fontsize=16, fontweight="bold")
# 1. Train loss.
# 1. Train + Val loss.
ax = axes[0]
sns.lineplot(data=train_df, x="epoch", y="total", ax=ax, marker="o", linewidth=2, color="crimson")
ax.set_title("Train Loss")
sns.lineplot(data=train_df, x="epoch", y="total", ax=ax, marker="o", linewidth=2, color="crimson", label="train")
if "loss" in val_df.columns:
sns.lineplot(data=val_df, x="epoch", y="loss", ax=ax, marker="s", linewidth=2, color="royalblue", label="val")
ax.set_title("Loss (train vs val)")
ax.set_xlabel("Epoch")
ax.set_ylabel("InfoNCE Loss")
ax.legend()
# 2. Val R@1 q→g.
ax = axes[1]