ShortcutCatcher/model.py at main · SmartData-Polito/ShortcutCatcher · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
import os
from autogluon.tabular import TabularPredictor, TabularDataset
import numpy as np

class Model:
    """
    Encapsulates training, prediction, and evaluation logic.
    Uses AutoGluon TabularPredictor internally.
    """

    def __init__(self, model_name='RandomForestGini'):
        self.model_name = model_name
        self.predictor = None

        if model_name == 'RandomForestGini':
            self.hyperparameters = {'RF': [{'criterion': 'gini', 'ag_args': {'name_suffix': 'Gini'}}]}
        elif model_name == 'NeuralNetTorch':
            self.hyperparameters = {'NN_TORCH': {}}
        elif model_name == 'TABICL':
            self.hyperparameters = {'TABICL': {}}
        elif model_name == 'XGBoost':
            self.hyperparameters = {'XGB': {}}
        elif model_name == 'FASTAI':
            self.hyperparameters = {'FASTAI': {}}
        elif model_name == 'REALMLP':
            self.hyperparameters = {'REALMLP': {}}
        else:
            raise ValueError(f"Unsupported model: {model_name}")

    def run(self, train_df, val_df, save_path, num_gpus=1):
        os.makedirs(save_path, exist_ok=True)
        train_data = TabularDataset(train_df)
        val_data = TabularDataset(val_df) if val_df is not None else None

        print(f"[Training] Start training {self.model_name}")
        self.predictor = TabularPredictor(label='class', path=save_path, verbosity=0).fit(
            train_data=train_data,
            tuning_data=val_data,
            num_gpus=num_gpus,
            hyperparameters=self.hyperparameters,
            fit_weighted_ensemble=False,
        )
        print(f"[Training] Model saved at {save_path}")
        return self.predictor

    def evaluate(self, data, name="", n_bins=15):
        print(f"[Eval] Evaluating on {name} ...")

        # 1) leaderboard（原有）
        df = self.predictor.leaderboard(
            data,
            extra_metrics=[
                'f1_macro','f1_micro','precision_macro','precision_micro',
                'recall_macro','recall_micro'
            ]
        )

        # 2) 取 y_true（兼容 AutoGluon 的 label 名称）
        label_col = getattr(self.predictor, 'label', 'class')
        y_true = data[label_col].values

        # 3) 预测概率（DataFrame，列顺序即类别顺序）
        proba_df = self.predictor.predict_proba(data)  # DataFrame
        proba_values = proba_df.values

        # 4) 计算 ECE
        ece = self.compute_ece_from_proba(y_true, proba_values, n_bins=n_bins)

        # 5) 并入 leaderboard（同一数据上 ECE 是单个标量，这里复制到每行，便于保存/对齐）
        df['ece'] = ece

        print(f"[Eval-{name}] Accuracy = {df['score_test'].values[0]:.4f} | ECE = {ece:.4f}")
        return df


    def compute_ece_from_proba(self, y_true_idx, proba_values, n_bins=15):
        conf = proba_values.max(axis=1)
        preds_idx = proba_values.argmax(axis=1)
        correct = (preds_idx == y_true_idx).astype(float)

        bins = np.linspace(0, 1, n_bins + 1)
        ece, N = 0.0, len(y_true_idx)
        for i in range(n_bins):
            lo, hi = bins[i], bins[i + 1]
            mask = (conf > lo) & (conf <= hi) if i > 0 else (conf >= lo) & (conf <= hi)
            if not np.any(mask):
                continue
            acc_bin  = correct[mask].mean()
            conf_bin = conf[mask].mean()
            ece += (mask.sum() / N) * abs(acc_bin - conf_bin)
        return float(ece)