Aim 实验追踪 — 开源、本地优先的高性能 ML 实验可视化工具¶

# 安装 Aim
pip install aim                            # 基础安装（含 UI 服务器）

# 初始化实验仓库（类似 git init）
mkdir my_experiments && cd my_experiments
aim init                                   # 创建 .aim 目录（存储所有实验数据）

# 启动可视化界面
aim up                                     # 默认在 http://localhost:43800 打开
aim up --port 8080                         # 自定义端口

# 验证安装
python -c "import aim; print(aim.__version__)"

from aim import Run                                     # Run 是 Aim 的核心对象

# 创建一个实验 Run（类似开始一个实验记录本）
run = Run(
    repo="/path/to/my_experiments",        # 实验仓库路径（默认当前目录）
    experiment="T2D菌群分类实验",           # 实验名（用于分组）
)

# 记录超参数（支持任意嵌套字典）
run["hparams"] = {                                      # 直接赋值即可
    "model": "RandomForest",
    "n_estimators": 100,
    "data": {
        "n_samples": 1000,
        "n_features": 50,
        "train_ratio": 0.8,
    }
}

# 记录指标（支持任意命名，实时更新到 UI）
for epoch in range(100):
    train_loss = 0.5 - epoch * 0.003
    val_auc = 0.75 + epoch * 0.002

    run.track(train_loss,                               # 要记录的值
              name="loss",                              # 指标名称
              step=epoch,                               # X 轴步数
              context={"subset": "train"})              # 上下文（用于分组）

    run.track(val_auc,
              name="auc",
              step=epoch,
              context={"subset": "val"})

# 记录图像（自动显示在 UI 中）
import numpy as np
from aim import Image
image_array = np.random.rand(64, 64, 3)               # 示例图片
run.track(Image(image_array), name="sample_image", step=0)

# 记录文本（例如模型预测结果）
from aim import Text
run.track(Text("预测：健康，概率：0.85"), name="predictions", step=0)

# 关闭 Run（完成记录）
run.close()

# Aim 与各训练框架的集成

# 1. 与 PyTorch 训练集成
from aim import Run
from aim.pytorch import track_gradients_dists, track_params_dists  # 追踪梯度分布

run = Run()
run["hparams"] = {"lr": 0.001, "batch_size": 32}

# 在训练循环中
for step, batch in enumerate(dataloader):
    loss = model(batch)
    optimizer.zero_grad()
    loss.backward()

    # 可选：追踪梯度分布（帮助发现梯度消失/爆炸）
    track_gradients_dists(run, model, step=step)

    optimizer.step()
    run.track(loss.item(), name="loss", step=step, context={"subset": "train"})

# 2. 与 HuggingFace Trainer 集成
from transformers import TrainingArguments

training_args = TrainingArguments(
    output_dir="./output",
    report_to="aim",                       # 只需设置这一行！
    # ... 其他参数
)

# 3. 与 PyTorch Lightning 集成
from aim.pytorch_lightning import AimLogger

aim_logger = AimLogger(
    experiment="my_experiment",
    repo="./aim_logs",                     # 存储路径
)

trainer = pl.Trainer(
    logger=aim_logger,
    # ... 其他参数
)

# 批量比较多组超参数实验（Aim 最擅长的场景）

from aim import Run
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import cross_val_score
import numpy as np

# 要比较的超参数组合
configs = [
    {"n_estimators": 50, "max_depth": 5},
    {"n_estimators": 100, "max_depth": 10},
    {"n_estimators": 200, "max_depth": 15},
    {"n_estimators": 100, "max_depth": None},  # None = 不限深度
]

X = np.random.randn(200, 30)               # 模拟数据
y = np.random.randint(0, 2, 200)

for config in configs:
    run = Run(experiment="RF超参数搜索")    # 每次实验创建一个 Run
    run["hparams"] = config               # 记录这次实验的超参数

    rf = RandomForestClassifier(**config, random_state=42)
    cv_scores = cross_val_score(rf, X, y, cv=5, scoring="roc_auc")

    # 记录 5 折交叉验证每折的结果
    for fold_idx, score in enumerate(cv_scores):
        run.track(score, name="cv_auc", step=fold_idx,
                  context={"fold": fold_idx})

    run.track(cv_scores.mean(), name="mean_auc", step=0)
    run.track(cv_scores.std(), name="std_auc", step=0)

    print(f"config={config}, AUC={cv_scores.mean():.4f}")
    run.close()

# 启动 Aim UI 对比所有实验
# aim up
# 在界面里可以把所有 Run 的 mean_auc 放在一起对比