/
scale_configs.py
55 lines (52 loc) · 1.38 KB
/
scale_configs.py
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
SCALE_CONFIGS = {
"debug": {
'batch_size': 1024,
'learning_rate': 1e-4,
'train_num_samples': 128_000,
'warmup': 500,
'model': 'ViT-B-32',
'beta2': None
},
"small": {
'batch_size': 4096,
'learning_rate': 5e-4,
'train_num_samples': 12_800_000,
'warmup': 500,
'model': 'ViT-B-32',
'beta2': None,
},
"medium": {
'batch_size': 4096,
'learning_rate': 5e-4,
'train_num_samples': 128_000_000,
'warmup': 500,
'model': 'ViT-B-32',
'beta2': None
},
"large": {
'batch_size': 8192,
'learning_rate': 5e-4,
'train_num_samples': 1_280_000_000,
'warmup': 500,
'model': 'ViT-B-16',
'beta2': None
},
"xlarge": {
'batch_size': 90112,
'learning_rate': 1e-3,
'train_num_samples': 12_800_000_000,
'warmup': 10000,
'model': 'ViT-L-14',
'beta2': 0.95
}
}
SIMPLE_NAMES = ['debug', 'small', 'medium', 'large', 'xlarge']
def available_scales(simple_names = False):
if simple_names:
return SIMPLE_NAMES
else:
return sorted(list(SCALE_CONFIGS.keys()))
def get_scale_config(scale):
if scale not in SCALE_CONFIGS:
raise ValueError(f'Unknown scale: {scale}. Please use one of {available_scales()}')
return SCALE_CONFIGS[scale]