-
Notifications
You must be signed in to change notification settings - Fork 13
/
Copy pathrun.py
181 lines (161 loc) · 7.07 KB
/
run.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
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
import os
import torch
import numpy as np
import symbolicregression
import requests
from symbolicregression.envs import build_env
from symbolicregression.model import build_modules
from parsers import get_parser
from symbolicregression.trainer import Trainer
from symbolicregression.slurm import init_signal_handler, init_distributed_mode
from evaluate import evaluate_pmlb, evaluate_pmlb_mcts, evaluate_in_domain, evaluate_mcts_in_domain
if __name__ == '__main__':
#load E2E model
model_path = "./symbolicregression/saved_models/model.pt"
try:
if not os.path.isfile(model_path):
url = "https://dl.fbaipublicfiles.com/symbolicregression/model1.pt"
r = requests.get(url, allow_redirects=True)
open(model_path, 'wb').write(r.content)
if not torch.cuda.is_available():
model = torch.load(model_path, map_location=torch.device('cpu'))
else:
model = torch.load(model_path)
print("Model successfully loaded!")
print(model.embedder)
except Exception as e:
print("ERROR: model not loaded! path was: {}".format(model_path))
print(e)
#load data:
parser = get_parser()
params = parser.parse_args()
init_distributed_mode(params)
if params.is_slurm_job:
init_signal_handler()
np.random.seed(params.seed)
torch.manual_seed(params.seed)
torch.cuda.manual_seed(params.seed)
# CPU / CUDA
if not params.cpu:
assert torch.cuda.is_available()
params.eval_only = True
symbolicregression.utils.CUDA = not params.cpu
env = build_env(params)
env.rng = np.random.RandomState(0)
modules = build_modules(env, params)
trainer = Trainer(modules, env, params)
params.batch_size_eval = 1
params.multi_gpu = False
params.is_slurm_job = False
params.random_state = 14423
params.eval_verbose_print = True
params.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(params.no_prefix_cache, params.no_seq_cache)
if (params.no_prefix_cache==False) and (params.no_seq_cache == False):
params.cache = 'b' #both cachings are arctivated
elif (params.no_prefix_cache==True) and (params.no_seq_cache ==False):
params.cache = 's' #only sequence caching
elif (params.no_prefix_cache==False) and (params.no_seq_cache==True):
params.cache = 'k' #only top-k caching
else:
params.cache = 'n' # no caching
# evaluate functions
if params.eval_in_domain:
scores = evaluate_in_domain(
trainer,
params,
model,
"valid1",
"functions",
verbose=True,
ablation_to_keep=None,
save=True,
logger=None,
save_file=params.save_eval_dic,)
print("Pare-trained E2E scores: ", scores)
if params.eval_mcts_in_domain:
scores = evaluate_mcts_in_domain(
trainer,
params,
model,
"valid1",
"functions",
verbose=True,
ablation_to_keep=None,
save=True,
logger=None,
save_file=None,
save_suffix="./eval_result/eval_in-domain_tpsr_l{}_b{}_k{}_r{}_cache{}_noise{}.csv".format(params.lam,
params.num_beams,
params.width,
params.rollout,
params.cache,
params.eval_noise_gamma),
)
print("TPSR scores: ", scores)
if params.eval_on_pmlb:
target_noise = params.target_noise
random_state = params.random_state
data_type = params.pmlb_data_type
save = params.save_results
if data_type == "feynman":
filter_fn = lambda x: x["dataset"].str.contains("feynman")
elif data_type == "strogatz":
print("Strogatz data")
filter_fn = lambda x: x["dataset"].str.contains("strogatz")
elif data_type == "603_fri_c0_250_50":
filter_fn = lambda x: x["dataset"].str.contains("603_fri_c0_250_50")
else:
filter_fn = lambda x: ~(
x["dataset"].str.contains("strogatz")
| x["dataset"].str.contains("feynman")
)
pmlb_scores = evaluate_pmlb(
trainer,
params,
model,
target_noise=target_noise,
verbose=params.eval_verbose_print,
random_state=random_state,
save=save,
filter_fn=filter_fn,
save_file=None,
save_suffix="./eval_result/eval_pmlb_pretrained_{}.csv".format(params.beam_type),
)
print("Pre-trained E2E scores: ", pmlb_scores)
if params.eval_mcts_on_pmlb:
target_noise = params.target_noise
random_state = params.random_state
data_type = params.pmlb_data_type
save = params.save_results
if data_type == "feynman":
filter_fn = lambda x: x["dataset"].str.contains("feynman")
elif data_type == "strogatz":
print("Strogatz data")
filter_fn = lambda x: x["dataset"].str.contains("strogatz")
elif data_type == "603_fri_c0_250_50":
filter_fn = lambda x: x["dataset"].str.contains("603_fri_c0_250_50")
else:
filter_fn = lambda x: ~(
x["dataset"].str.contains("strogatz")
| x["dataset"].str.contains("feynman")
)
pmlb_scores = evaluate_pmlb_mcts(
trainer,
params,
model,
target_noise=target_noise,
verbose=params.eval_verbose_print,
random_state=random_state,
save=save,
filter_fn=filter_fn,
save_file=None,
save_suffix="./eval_result/eval_{}_tpsr_l{}_b{}_k{}_r{}_cache{}_noise{}.csv".format(params.pmlb_data_type,
params.lam,
params.num_beams,
params.width,
params.rollout,
params.cache,
params.target_noise),
)
print("TPSR scores: ", pmlb_scores)