-
Notifications
You must be signed in to change notification settings - Fork 0
/
configs.py
102 lines (101 loc) · 2.35 KB
/
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
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
92
93
94
95
96
97
98
99
100
101
102
max_pow = 5
aquamam_toy = {
"model_args": {
"toy_args": {"is_toy": True, "max_pow": max_pow},
"L": 6,
"d_model": 64,
"nhead": 8,
"dropout": 0.0,
"num_layers": 3,
"n_bins": 50257, # GPT-3 vocaublary size.
},
"batch_size": 128,
"patience": 5,
"num_workers": 2,
"lr": 1e-4,
"epochs": 1000,
"test_batch_size": 1536,
"beam_k": 1,
}
aquamam_mog_toy = {
"model_args": {
"toy_args": {"is_toy": True, "max_pow": max_pow},
"L": 6,
"d_model": 64,
"nhead": 8,
"dropout": 0.0,
"num_layers": 3,
"n_comps": 512,
},
"batch_size": 128,
"patience": 5,
"num_workers": 2,
"lr": 1e-4,
"epochs": 1000,
"test_batch_size": 1536,
"beam_k": 1,
}
aquamam_solid = {
"model_args": {
"toy_args": {"is_toy": False},
"L": 6,
"d_model": 512,
"nhead": 8,
"dropout": 0.1,
"num_layers": 6,
"n_bins": 500,
},
"batch_size": 128,
"patience": 5,
"num_workers": 2,
"lr": 1e-4,
"epochs": 1000,
"test_batch_size": 1024,
"predict_batch_size": 1536,
"beam_k": 1,
}
ipdf_toy = {
"model_args": {
"toy_args": {
"is_toy": True,
"max_pow": max_pow,
"visual_embedding_size": 2048,
},
"resnet": "resnet50",
"L": 3,
"n_hidden_nodes": 256,
"mlp_layers": 4,
},
"neg_samples": 4095,
"batch_size": 128,
"num_workers": 2,
"lr": 1e-4,
"warmup_steps": 1000,
"iterations": 300000,
"number_queries": 2000000,
"test_batch_size": 1,
}
# See: https://github.com/google-research/google-research/tree/master/implicit_pdf#reproducing-symsol-results
# and Section S8.
ipdf_solid = {
"model_args": {
"toy_args": {"is_toy": False},
"resnet": "resnet50",
"L": 3,
"n_hidden_nodes": 256,
"mlp_layers": 4,
},
"neg_samples": 4095,
"batch_size": 128,
"num_workers": 2,
"lr": 1e-4,
"warmup_steps": 1000,
"iterations": 300000,
"number_queries": 2000000,
"test_batch_size": 1,
}
configs = {
"aquamam": {"toy": aquamam_toy, "cube": aquamam_solid, "cylinder": aquamam_solid},
"ipdf": {"toy": ipdf_toy, "cube": ipdf_solid, "cylinder": ipdf_solid},
"aquamam_mog": {"toy": aquamam_mog_toy},
}