HyperCLIP/scale_configs.py at main · locuslab/HyperCLIP · 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
"""Training recipes used by the HyperCLIP paper."""

def _recipe(model, samples=128_000_000):
    return {
        "batch_size": 1500,
        "learning_rate": 5e-4,
        "train_num_samples": samples,
        "warmup": 500,
        "model": model,
        "beta2": None,
    }


SCALE_CONFIGS = {
    # Eight matched backbone pairs from the main experiments.
    "rep_b0_siglip": _recipe("TimEfNetb0"),
    "rep_b0_hyperclip": _recipe("HypTimEfNetb0"),
    "rep_b1_siglip": _recipe("TimEfNetb1"),
    "rep_b1_hyperclip": _recipe("HypTimEfNetb1"),
    "rep_b2_siglip": _recipe("TimEfNet"),
    "rep_b2_hyperclip": _recipe("HypTimEfNet"),
    "rep_m0_siglip": _recipe("TimMobNet"),
    "rep_m0_hyperclip": _recipe("HypTimMobNet"),
    "rep_m1_siglip": _recipe("TimMobNetSmall"),
    "rep_m1_hyperclip": _recipe("HypTimMobNetSmall"),
    "rep_t0_siglip": _recipe("TimTinyNet"),
    "rep_t0_hyperclip": _recipe("HypTimTinyNet"),
    "rep_e0_siglip": _recipe("TimEdgNetSmall"),
    "rep_e0_hyperclip": _recipe("HypTimEdgNetSmall"),
    "rep_v0_siglip": _recipe("MobViT"),
    "rep_v0_hyperclip": _recipe("HypMobViT"),

    # Reduced-budget, matched EfficientNet-B0 controls.
    "rep_b0_base32": _recipe("HypTimEfNetb0", 32_000_000),
    "rep_b0_token32": _recipe("HypTimEfNetb0_token", 32_000_000),
    "rep_b0_linear32": _recipe("HypTimEfNetb0_linear", 32_000_000),
    "rep_b0_conv32": _recipe("HypTimEfNetb0_conv", 32_000_000),
}


def available_scales():
    return sorted(SCALE_CONFIGS)


def get_scale_config(scale):
    try:
        return SCALE_CONFIGS[scale]
    except KeyError as error:
        raise ValueError(
            f"Unknown scale {scale!r}; choose one of {available_scales()}"
        ) from error