-
Notifications
You must be signed in to change notification settings - Fork 8
/
run_lfads_multi.py
154 lines (127 loc) · 7.27 KB
/
run_lfads_multi.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
import argparse
import os
import pickle
import torch
import torchvision
import torchvision.transforms as trf
import torch.optim as opt
from trainer import RunManager
from scheduler import LFADS_Scheduler
from objective import LFADS_Loss, LogLikelihoodPoisson
from lfads import LFADS_MultiSession_Net
from utils import read_data, load_parameters
from plotter import Plotter
from dataset import LFADS_MultiSession_Dataset, SessionLoader
parser = argparse.ArgumentParser()
parser.add_argument('-d', '--data_path', type=str)
parser.add_argument('-p', '--hyperparameter_path', type=str)
parser.add_argument('-o', '--output_dir', default='/tmp', type=str)
parser.add_argument('--max_epochs', default=2000, type=int)
parser.add_argument('--batch_size', default=None, type=int)
parser.add_argument('-t', '--use_tensorboard', action='store_true', default=False)
parser.add_argument('-r', '--restart', action='store_true', default=False)
parser.add_argument('-c', '--do_health_check', action='store_true', default=False)
def main():
args = parser.parse_args()
device = 'cuda' if torch.cuda.is_available() else 'cpu'
hyperparams = load_parameters(args.hyperparameter_path)
data_name = os.path.splitext(args.data_path.split('/')[-1])[0]
model_name = hyperparams['model_name']
mhp_list = [key.replace('size', '').replace('_', '')[:4] + str(val) for key, val in hyperparams['model'].items() if 'size' in key]
mhp_list.sort()
hyperparams['run_name'] = '_'.join(mhp_list)
save_loc = '%s/%s/%s/%s/'%(args.output_dir, data_name, model_name, hyperparams['run_name'])
if not os.path.exists(save_loc):
os.makedirs(save_loc)
data_dict = pickle.load(open(args.data_path, 'rb'))
keys = [key for key in data_dict.keys() if type(data_dict[key]) is dict]
keys.sort()
train_data_list = [data_dict[key]['train_data'] for key in keys]
valid_data_list = [data_dict[key]['valid_data'] for key in keys]
train_ds = LFADS_MultiSession_Dataset(train_data_list, device=device)
valid_ds = LFADS_MultiSession_Dataset(valid_data_list, device=device)
train_dl = SessionLoader(train_ds, sampler=torch.utils.data.RandomSampler(train_ds))
valid_dl = SessionLoader(valid_ds)
transforms = trf.Compose([])
loglikelihood = LogLikelihoodPoisson(dt=float(data_dict['dt']))
objective = LFADS_Loss(loglikelihood = loglikelihood,
loss_weight_dict = {'kl': hyperparams['objective']['kl'],
'l2': hyperparams['objective']['l2']},
l2_con_scale = hyperparams['objective']['l2_con_scale'],
l2_gen_scale = hyperparams['objective']['l2_gen_scale']).to(device)
W_in_list = [torch.Tensor(data_dict[key]['W_in']).to(device) for key in keys]
W_out_list = [torch.Tensor(data_dict[key]['W_out']).to(device) for key in keys]
b_in_list = [torch.Tensor(data_dict[key]['b_in']).to(device) for key in keys]
b_out_list = [torch.Tensor(data_dict[key]['b_out']).to(device) for key in keys]
model = LFADS_MultiSession_Net(W_in_list = W_in_list,
W_out_list = W_out_list,
b_in_list = b_in_list,
b_out_list = b_out_list,
factor_size = hyperparams['model']['factor_size'],
g_encoder_size = hyperparams['model']['g_encoder_size'],
c_encoder_size = hyperparams['model']['c_encoder_size'],
g_latent_size = hyperparams['model']['g_latent_size'],
u_latent_size = hyperparams['model']['u_latent_size'],
controller_size = hyperparams['model']['c_controller_size'],
generator_size = hyperparams['model']['generator_size'],
prior = hyperparams['model']['prior'],
clip_val = hyperparams['model']['clip_val'],
dropout = hyperparams['model']['dropout'],
do_normalize_factors = hyperparams['model']['normalize_factors'],
max_norm = hyperparams['model']['max_norm'],
device = device).to(device)
total_params = 0
for ix, (name, param) in enumerate(model.named_parameters()):
print(ix, name, list(param.shape), param.numel(), param.requires_grad)
total_params += param.numel()
print('Total parameters: %i'%total_params)
optimizer = opt.Adam(model.parameters(),
lr=hyperparams['optimizer']['lr_init'],
betas=hyperparams['optimizer']['betas'],
eps=hyperparams['optimizer']['eps'])
scheduler = LFADS_Scheduler(optimizer = optimizer,
mode = 'min',
factor = hyperparams['scheduler']['scheduler_factor'],
patience = hyperparams['scheduler']['scheduler_patience'],
verbose = True,
threshold = 1e-4,
threshold_mode = 'abs',
cooldown = hyperparams['scheduler']['scheduler_cooldown'],
min_lr = hyperparams['scheduler']['lr_min'])
num_steps = train_data_list[0].shape[1]
TIME = torch._np.arange(0, num_steps*data_dict['dt'], data_dict['dt'])
plotter = {'train' : Plotter(time=TIME),
'valid' : Plotter(time=TIME)}
if args.use_tensorboard:
import importlib
if importlib.util.find_spec('torch.utils.tensorboard'):
tb_folder = save_loc + 'tensorboard/'
if not os.path.exists(tb_folder):
os.mkdir(tb_folder)
elif os.path.exists(tb_folder) and args.restart:
os.system('rm -rf %s'%tb_folder)
os.mkdir(tb_folder)
from torch.utils.tensorboard import SummaryWriter
writer = SummaryWriter(tb_folder)
rm_plotter = plotter
else:
writer = None
rm_plotter = None
else:
writer = None
rm_plotter = None
run_manager = RunManager(model = model,
objective = objective,
optimizer = optimizer,
scheduler = scheduler,
train_dl = train_dl,
valid_dl = valid_dl,
transforms = transforms,
writer = writer,
plotter = rm_plotter,
max_epochs = args.max_epochs,
save_loc = save_loc,
do_health_check = args.do_health_check)
run_manager.run()
if __name__ == '__main__':
main()