LogisticCalibrationBenchmark/benchmark_binary.py at main · eugeneberta/LogisticCalibrationBenchmark · 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
import torch
import numpy as np
import pandas as pd
from tqdm import tqdm
from utils import test_calibrator
from tabrepo import load_repository
from probmetrics.metrics import Metrics
from probmetrics.calibrators import get_calibrator

def prepare_dataset(repo, dataset, fold, config, cache):
    """Load and subsample dataset splits once, reuse for all calibrators."""
    key = (dataset, fold, config)
    if key in cache:
        return cache[key]

    metrics = Metrics.from_names(['logloss', 'brier', 'accuracy', 'ece-15'])

    p_cal = repo.predict_val(dataset=dataset, fold=fold, config=config, binary_as_multiclass=False)
    y_cal = repo.labels_val(dataset=dataset, fold=fold)
    p_test = repo.predict_test(dataset=dataset, fold=fold, config=config, binary_as_multiclass=True)
    y_test = repo.labels_test(dataset=dataset, fold=fold)

    n_cal = len(p_cal)
    n_test = len(p_test)

    # Subsampling large datasets:
    if n_cal >= 10000:
        np.random.seed(123)
        idx = np.arange(0, n_cal)
        rand_idx = np.random.choice(idx, 10000, replace=False)
        p_cal = p_cal[rand_idx]
        y_cal = y_cal[rand_idx]
        n_cal = len(p_cal)

    base_metrics = metrics.compute_all_from_labels_probs(
        torch.as_tensor(y_test), torch.as_tensor(p_test)
    )

    base_results = {
        'dataset': dataset,
        'fold': fold,
        'config': config,
        'cal_size': n_cal,
        'test_size': n_test,
    }
    base_results.update({f'base_{key}': value.item() for key, value in base_metrics.items()})

    cache[key] = (p_cal, y_cal, p_test, y_test, base_results, metrics)
    return cache[key]


if __name__ == "__main__":
    repo = load_repository("D244_F3_C1530_200")
    configs = pd.read_csv('results/binary/configs.csv')

    calibrator_factories = {
        'iso': lambda: get_calibrator('isotonic'),
        'guo_ts': lambda: get_calibrator('guo-ts'),
        'prob_ts': lambda: get_calibrator('ts-mix'),
        'ts': lambda: get_calibrator('linear-scaling'),
        'as': lambda: get_calibrator('affine-scaling'),
        'qs': lambda: get_calibrator('quadratic-scaling'),
    }

    data_cache = {}
    aggregated_results = {}

    for cal_name, factory in calibrator_factories.items():
        print(f"\n=== Running benchmark for calibrator: {cal_name} ===")
        for _, row in tqdm(configs.iterrows(), total=len(configs)):
            dataset, fold, config = row['dataset'], row['fold'], row['tuned_config']

            p_cal, y_cal, p_test, y_test, base_results, metrics = prepare_dataset(
                repo, dataset, fold, config, data_cache
            )

            key = (dataset, fold, config)
            if key not in aggregated_results:
                aggregated_results[key] = base_results.copy()

            cal = factory()
            aggregated_results[key] = test_calibrator(
                cal, cal_name, metrics, aggregated_results[key], p_cal, y_cal, p_test, y_test
            )

    all_results = list(aggregated_results.values())
    df = pd.DataFrame(all_results)
    df.to_csv('results/binary/results.csv', index=False)