-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathtest.py
184 lines (151 loc) · 7.33 KB
/
test.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
182
183
184
import numpy as np
import pickle
import torch
from pathlib import Path
from rich.panel import Panel
from rich.pretty import Pretty
from rich.console import Console
from rich.progress import Progress, SpinnerColumn, TimeElapsedColumn
import time as tt
from utils.utils import plot_source_and_target_mnist, binary, save_gif_frame_mnist
from utils.data_utils import gen_mnist_data, reverse_normalize_dataset, normalize_dataset_with_metadata
from utils.model_utils import get_model_before_after
import argparse
def check_model_task(args):
if args.task == 'gaussian2mnist':
assert args.model in ['tunet++', 'aunet', 'unet++', 'unet']
args.time_expand = False
else:
assert args.model in ['mlp', 'unet++', 'unet']
args.time_expand = True
def main():
parser = argparse.ArgumentParser(description='Train SDE')
parser.add_argument('--task', type=str, default='gaussian2mnist', required=True)
parser.add_argument('--model', type=str, default='unet++', required=True)
parser.add_argument('--checkpoint', type=str, default=None, required=True)
parser.add_argument('--seed', type=int, default=233)
parser.add_argument('--change_epsilons', action='store_true')
parser.add_argument('--scheduler', type=str, default=None)
parser.add_argument('--lr', type=float, default=1e-4)
parser.add_argument('--iter_nums', type=int, default=1)
parser.add_argument('--epoch_nums', type=int, default=2)
parser.add_argument('-b', '--batch_size', type=int, default=8000)
parser.add_argument('-n','--normalize', action='store_true')
parser.add_argument('--tarined_data', action='store_true')
parser.add_argument('--filter_number', type=int)
parser.add_argument('--debug', action='store_true')
args = parser.parse_args()
check_model_task(args)
seed = args.seed
torch.manual_seed(seed)
if torch.cuda.is_available():
torch.cuda.manual_seed_all(seed)
np.random.seed(seed)
experiment_name = args.task
if args.change_epsilons:
experiment_name += '_change_epsilons'
if args.filter_number is not None and 'mnist' in args.task:
experiment_name += f'_filter{args.filter_number}'
if args.debug:
log_dir = Path('experiments') / 'debug' / 'test' / tt.strftime("%Y-%m-%d/%H_%M_%S/")
else:
log_dir = Path('experiments') / experiment_name / 'test' / tt.strftime("%Y-%m-%d/%H_%M_%S/")
ds_cached_dir = Path('experiments') / experiment_name / 'data'
log_dir.mkdir(parents=True, exist_ok=True)
ds_cached_dir.mkdir(parents=True, exist_ok=True)
args.log_dir = log_dir
args.ds_cached_dir = ds_cached_dir
if args.task == 'gaussian2mnist':
args.dim = 1
else:
args.dim = 2
args.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
main_worker(args)
def main_worker(args):
console = Console(record=True, color_system='truecolor')
pretty = Pretty(args.__dict__, expand_all=True)
panel = Panel(pretty, title='Arguments', expand=False, highlight=True)
console.log(panel)
console.log(f"Saving to {Path.absolute(args.log_dir)}")
model, before_train, after_train = get_model_before_after(args)
# console.log(model)
if args.checkpoint is not None:
try:
model.load_state_dict(torch.load(args.checkpoint))
except:
console.log(":warning-emoji: [bold red blink] load checkpoint failed [/]\ncheckpoint: {}\nExit".format(args.checkpoint))
return None
console.log("load checkpoint from {}".format(args.checkpoint))
console.log(f"Model {model.__class__.__name__} Parameters: {int(sum(p.numel() for p in model.parameters())/1e6)}M")
real_metadata = pickle.loads(open(args.ds_cached_dir / 'real_mean_std.pkl', 'rb').read())
if args.tarined_data:
ds_cached_files = [f for f in args.ds_cached_dir.iterdir() if f.name.startswith('new_ds_')]
temp_ds = pickle.loads(open(ds_cached_files[0], 'rb').read())
data = {
"ts": temp_ds.ts,
"bridge": temp_ds.bridge,
"drift": temp_ds.drift,
"source": temp_ds.source,
}
data = reverse_normalize_dataset(temp_ds.metadata, **data)
# data = normalize_dataset_with_metadata(real_metadata, **data)
test_ts = data['ts']
test_bridge = data['bridge']
test_drift = data['drift']
test_source = data['source']
else:
test_ts, test_bridge, test_drift, test_source, _ = gen_mnist_data(nums=1000)
pred_bridge = torch.zeros_like(test_bridge)
pred_drift = torch.zeros_like(test_drift)
pred_bridge[0, :] = test_source
# model.eval()
sigma=1
console.rule("[bold deep_sky_blue1 blink]Testing")
with Progress(
SpinnerColumn(spinner_name='moon'),
*Progress.get_default_columns(),
TimeElapsedColumn(),
transient=False,
) as progress:
task1 = progress.add_task("[gold]Predicting", total=len(test_ts) - 1)
with torch.no_grad():
for i in range(len(test_ts) - 1):
dt = (test_ts[i+1] - test_ts[i])
test_source_reshaped = test_source
if args.time_expand:
test_ts_reshaped = test_ts[i].repeat(test_source.shape[0]).reshape(-1, 1, 1, 1).repeat(1, 1, 28, 28)
else:
test_ts_reshaped = test_ts[i].repeat(test_source.shape[0])
# test_ts_reshaped = torch.unsqueeze(test_ts[i].repeat(test_source.shape[0]), dim=0).T
pred_bridge_reshaped = pred_bridge[i]
ret = normalize_dataset_with_metadata(real_metadata, source=test_source_reshaped, ts=test_ts_reshaped, bridge=pred_bridge_reshaped)
test_ts_reshaped = ret['ts']
pred_bridge_reshaped = ret['bridge']
test_source_reshaped = ret['source']
if args.time_expand:
x = torch.concat([test_source_reshaped, test_ts_reshaped, pred_bridge_reshaped], axis=1)
time = None
else:
x = torch.concat([test_source_reshaped, pred_bridge_reshaped], axis=1)
time = test_ts_reshaped.to(args.device)
if before_train is not None:
x = before_train(x)
x = x.to(args.device)
model = model.to(args.device)
if args.debug:
print(x.shape, time.shape if time is not None else None)
dydt = model(x, time) if time is not None else model(x)
dydt = dydt.cpu()
if after_train is not None:
dydt = after_train(dydt)
pred_drift[i]=dydt
dydt = reverse_normalize_dataset(real_metadata, bridge=dydt)['bridge']
diffusion = sigma * torch.sqrt(dt) * torch.randn_like(dydt)
pred_bridge[i+1] = pred_bridge[i] + dydt * dt + diffusion[:]
progress.update(task1, advance=1)
plot_source_and_target_mnist(test_bridge[-1, :25], binary(pred_bridge[-1, :25]), save_path=args.log_dir / 'source_and_target.jpg')
save_gif_frame_mnist(pred_bridge[:, :25], args.log_dir, 'pred.gif')
save_gif_frame_mnist(pred_bridge[:, :25], args.log_dir, 'pred_norm.gif', norm=True)
if __name__ == '__main__':
main()
pass