-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.py
274 lines (234 loc) · 12.1 KB
/
main.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
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
import sys
import os
import json
import logging
import pickle
from pytorch_lightning.loggers import WandbLogger
from pytorch_lightning.callbacks import ModelCheckpoint
import configargparse
root_dir = "/nobackup/users/yankeson/Astronomy"
sys.path.insert(0, f"{root_dir}/ppae/")
from utils import *
from autoencoder import *
from dataset import *
device = 'cuda'
if __name__ == "__main__":
p = configargparse.ArgumentParser()
################# All Parameters ####################
# Very important ones (values required)
p.add_argument('--model_name', type=str, required=True,
help='Name of the model. Relevant for saving and plotting')
p.add_argument('--num_epochs', type=int, required=True,
help='Number of epochs to train')
p.add_argument('--lam_TV', type=float, required=True,
help='Penalty for total variation')
p.add_argument('--model_type', type=str, required=True,
help='Type of the model. decoder: decoder only. lstm: lstm encoder. transformer: vanilla transformer encoder')
p.add_argument('--starting_epoch', type=int, required=True,
help='Which epoch to start training')
p.add_argument('--checkpoint_every', type=int, required=True,
help='How often to save checkpoints')
p.add_argument('--data_type', type=str, required=True,
help='Whether to use large or small dataset')
p.add_argument('--TV_type', type=str, required=True,
help='Which TV loss type to use')
# Important ones (but with default values)
p.add_argument('--finetune_checkpoint', type=str, default=None, required=False,
help='Whether to finetune a model with larger dataset with one from smaller')
p.add_argument('--latent_only', action='store_true',
help='Whether to only optimize latents')
p.add_argument('--discrete', action='store_true',
help='Whether to use discrete version of the decoder (not a neural field)')
p.add_argument('--thin_resnet', action='store_true',
help='Whether to use thin resnet for the decoder')
p.add_argument('--random_shift', action='store_true',
help='Whether to random shift the first event in the dataset')
p.add_argument('--B', type=int, default=32, required=False,
help='Batch size')
p.add_argument('--lr', type=float, default=0.001, required=False,
help='Learning rate')
p.add_argument('--num_freqs', type=int, default=12, required=False,
help='Number of frequencies for the positional encoding')
p.add_argument('--latent_size', type=int, default=64, required=False,
help='Dimensionality of latent space')
p.add_argument('--hidden_size', type=int, default=512, required=False,
help='hidden size for the decoder ResNet')
p.add_argument('--hidden_blocks', type=int, default=5, required=False,
help='number of hidden blocks for the decoder ResNet')
p.add_argument('--lam_latent', type=float, default=0.0, required=False,
help='Penalty for norm of latents')
p.add_argument('--d_encoder_model', type=int, default=48, required=False,
help='Dimensionality of transformer token')
p.add_argument('--n_head', type=int, default=4, required=False,
help='Number of transformer heads')
p.add_argument('--num_encoder_layers', type=int, default=1, required=False,
help='Number of encoder layers for transformer/LSTM')
p.add_argument('--dim_feedforward', type=int, default=512, required=False,
help='Number of feedforward neurons in transformer/LSTM')
p.add_argument('--clip_val', type=float, default=10.0, required=False,
help='Value to clip gradient')
# Less important ones
p.add_argument('--num_workers', type=int, default=4, required=False,
help='Number of workers')
p.add_argument('--E_bins', type=int, default=3, required=False,
help='Number of energy bins to discretize for')
p.add_argument('--resolution', type=int, default=2048, required=False,
help='Resolution for mesh')
p.add_argument('--plotting_nbins', type=int, default=100, required=False,
help='Scale to normalize times for plottng')
# p.add_argument('--plotting_E_index', type=int, default=1, required=False,
# help='Which energy bin to plot')
opt = p.parse_args()
os.chdir(root_dir)
folder_path = f'{root_dir}/experiments/{opt.model_name}'
if not os.path.exists(folder_path):
os.makedirs(folder_path)
# Create relevant files
with open(f'{folder_path}/arguments.json', 'w') as file:
json.dump(vars(opt), file)
#### Load data
# Load and deserialize the list from the file
if opt.data_type == 'large':
filename = f'{root_dir}/Chandra_data/large_eventfiles_lifetime28800.pkl'
elif opt.data_type == 'large_filter':
filename = f'{root_dir}/Chandra_data/large_eventfiles_filtered_lifetime28800.pkl'
elif opt.data_type == 'large_filtermore':
filename = f'{root_dir}/Chandra_data/large_eventfiles_filteredmore_lifetime28800.pkl'
else:
filename = f'{root_dir}/Chandra_data/small_eventfiles_lifetime43200.pkl'
if opt.random_shift:
filename = filename[:-4] + '_randomshift.pkl'
with open(filename, 'rb') as file:
data_lst = pickle.load(file)
# Load into dataset and dataloader
if opt.data_type == 'small':
t_scale = 43200
else:
t_scale = 28800
data = RealEventsDataset(data_lst,E_bins=opt.E_bins,t_scale=t_scale)
loader = DataLoader(data, batch_size=opt.B, shuffle=True, num_workers=opt.num_workers, collate_fn=padding_collate_fn)
##### Set up validation plotting
if opt.data_type == 'large':
plotting_inds = [17219, 34935, 36634, 54247, 88181, 88609, # flares
49551, 51095, 1970, 4424, 42866, 74778, # dips
71, 159, 304, 381] # other random ids
elif opt.data_type == 'large_filtermore':
plotting_inds = [709, # flares
64, 333, 363, 397, 412, 591, 592, # dips
71, 159, 304, 381, 118, 42, 513, 832] # other random ids
elif opt.data_type == 'large_filter':
# Lightly filtered
plotting_inds = [297, 1940, 6294, 11123, 11197, # flares
290, 558, 911, 4683, 5587, 4997, # dips
71, 159, 304, 381, 2024] # other random ids
else:
plotting_inds = [i for i in range(8)] + [j for j in range(600,608)]
# Define the test set
batch = [data[i] for i in plotting_inds]
batch = padding_collate_fn(batch)
################## Create, train and save the NN model
def find_id_by_name(project_name, run_name):
api = wandb.Api()
runs = api.runs(path=f"{api.default_entity}/{project_name}")
for run in runs:
if run.name == run_name:
return run.id
return None
# Search for the project by name and get its ID
run_id = find_id_by_name('ppad', f"{opt.model_name}_lr00001")
if run_id:
wandb_logger = WandbLogger(project='ppad', name=f"{opt.model_name}_lr00001", id=run_id, resume='allow')
else:
wandb_logger = WandbLogger(project='ppad', name=f"{opt.model_name}_lr00001")
clip_val = 10
encoding = PositionalEncoding(num_freqs=opt.num_freqs)
checkpoint_callback = ModelCheckpoint(
dirpath=folder_path,
filename='model_{epoch}', # Customize the checkpoint filename
save_top_k=-1, # Save all checkpoints
every_n_epochs=opt.checkpoint_every # Save a checkpoint every 10 epochs
)
trainer = pl.Trainer(max_epochs=opt.starting_epoch+opt.num_epochs,
accelerator=device,
devices=1,
plugins=[DisabledSLURMEnvironment(auto_requeue=False)],
logger=wandb_logger,
callbacks=checkpoint_callback,
gradient_clip_val=opt.clip_val)
# callbacks=[StochasticWeightAveraging(swa_lrs=1e-2)])
# accumulate_grad_batches=5)
# )
# precision="16-mixed")
if opt.starting_epoch == 0:
modelclass = DiscreteAutoEncoder if opt.discrete else AutoEncoder
model = modelclass(opt, encoding, latent_num=len(data), test_batch=batch)
if opt.finetune_checkpoint is not None:
assert opt.data_type == 'large'
filename = f'{root_dir}/Chandra_data/large_eventfiles_filtered_lifetime28800.pkl'
with open(filename, 'rb') as file:
small_data_lst = pickle.load(file)
model = load_from_less_latents(model, small_data_lst, data_lst, opt.finetune_checkpoint)
del small_data_lst
history = trainer.fit(model, loader)
else:
search_pattern = os.path.join(folder_path, f'model_{opt.starting_epoch}epochs*.ckpt')
matching_files = glob(search_pattern)
if len(matching_files) == 0:
raise FileNotFoundError(f"No checkpoint files starting with 'model_{opt.starting_epoch}epochs' found in {folder_path}")
checkpoint_file = matching_files[0]
modelclass = DiscreteAutoEncoder if opt.discrete else AutoEncoder
model = modelclass.load_from_checkpoint(checkpoint_file, opt=opt, encoding=encoding, latent_num=len(data), test_batch=batch)
try:
history = trainer.fit(model, loader, ckpt_path=checkpoint_file)
except Exception as e:
print('Cannot resume trainer normally, doing something else')
checkpoint = torch.load(checkpoint_file)
# Re-initialize the trainer
del trainer
trainer = pl.Trainer(max_epochs=opt.starting_epoch+opt.num_epochs,
accelerator=device,
devices=1,
plugins=[DisabledSLURMEnvironment(auto_requeue=False)],
logger=wandb_logger,
callbacks=checkpoint_callback,
gradient_clip_val=opt.clip_val)
trainer.fit_loop.epoch_progress.current.completed = checkpoint['epoch']
trainer.fit_loop.epoch_loop._batches_that_stepped = checkpoint['global_step']
del checkpoint
history = trainer.fit(model, loader)
lrstr = f"{opt.lr:.0e}"
trainer.save_checkpoint(f'{folder_path}/model_{opt.starting_epoch+opt.num_epochs}epochs_lr{lrstr}.ckpt')
################ Inference
# model.to(device)
# B_test = 16
# test_loader = DataLoader(data, batch_size=B_test, collate_fn=padding_collate_fn)
# outputs = []
# for idx, batch in enumerate(test_loader):
# batch = todevice(batch, device)
# outputs.append(todevice(model(batch),'cpu'))
# # Plot total rates
# if opt.data_type == 'large':
# Tmax = 28800
# else:
# Tmax = 43200
# plt.figure(figsize=(12,9))
# for i, total_index in enumerate(plotting_inds):
# batch_index = total_index // B_test
# if opt.data_type == 'large':
# # index = total_index % 2 # Kind of random!
# index = 0
# else:
# index = total_index % B_test
# batch = outputs[batch_index]
# mask = batch['mask'][index]
# times = batch['event_list'][index,mask,0] * t_scale / 3600
# # rates = batch['rates'][index,mask] * Tmax / t_scale / opt.plotting_nbins
# total_mask = batch['total_mask'][index]
# total_times = batch['total_list'][index,total_mask,0] * t_scale / 3600
# total_rates = batch['total_rates'][index,total_mask] * Tmax / t_scale / opt.plotting_nbins
# plt.subplot(4,4,i+1)
# plt.hist(times, bins = opt.plotting_nbins)
# plt.plot(total_times, torch.sum(total_rates,dim=-1))
# plt.suptitle('Fitted vs true total rates',size=20)
# plt.tight_layout()
# plt.savefig(f'{folder_path}/total_rates_{opt.starting_epoch+opt.num_epochs}epochs.png')