def log_kv_or_listkv(dic, prefix): for k, v in dic.items(): if isinstance(v, (list, tuple)): for i, val in enumerate(v): if isinstance(val, (float, int)): logger.logkv("{}/{}/{}".format(prefix, k, i), val) elif isinstance(v, (float, int)): logger.logkv("{}/{}".format(prefix, k), v)
def ecb(epoch): params = list(system.parameters()) vparams = parameters_to_vector(params) error = (vparams - true_params).norm().item() logger.logkv('test/log10_paramerror', np.log10(error)) return
def ecb(epoch): torch.save(sys.state_dict(), os.path.join(logger.get_dir(), 'ckpts', 'model_{}.th'.format(epoch))) y_pred = gen_ypred_model(sys, val_u, val_y) rms = compute_rms(val_y[:, 25:], y_pred[:, 25:], y_std) val_rmse = float(rms.mean()) logger.logkv('test/val_rmse', val_rmse) if val_rmse < tracker.best_val_rmse: tracker.best_val_rmse = val_rmse torch.save(sys.state_dict(), os.path.join(logger.get_dir(), 'ckpts', 'best_model.th')) return
def tcb(epoch): with torch.no_grad(): tgt_test_pr = system.step_derivs(t_test, x_test) error = float(torch.nn.functional.mse_loss(tgt_test_pr, tgt_test)) logger.logkv('test/log10_error', np.log10(error)) last_10_errors._arr.append(np.log10(error)) if len(last_10_errors._arr) > 100: last_10_errors._arr = last_10_errors._arr[-100:] l10err = torch.tensor(last_10_errors._arr) convcrit = float((l10err.min() - l10err.max()).abs()) logger.logkv('test/log10_convcrit', np.log10(convcrit)) if convcrit < 1e-3: return True return False
def tcb(epoch): params = list(system.parameters()) vparams = parameters_to_vector(params) error = (vparams - true_params).norm().item() last_10_errors._arr.append(float(error)) logger.logkv('test/log10_error', np.log10(error)) if len(last_10_errors._arr) > 10: last_10_errors._arr = last_10_errors._arr[-10:] l10err = torch.tensor(last_10_errors._arr) convcrit = float((l10err.min() - l10err.max()).abs()) logger.logkv('test/log10_convcrit', np.log10(convcrit)) if convcrit < 1e-4: return True return False
def ecb(k): logger.logkv('train/epoch', k) logger.logkv('train/elapsed_time', time() - start_time) # update lr learner_opt_kwargs['lr'] = lr_sched(k) logger.logkv('train/lr', lr_sched(k)) torch.save(system.state_dict(), os.path.join(logger.get_dir(), 'ckpts', 'best_model.th')) torch.save(xsm, os.path.join(logger.get_dir(), 'ckpts', 'best_xsm.th')) logger.logkv('train/mu', float(system.logmu.exp()))
def train(self, params, callbacks=[]): vparams0 = parameters_to_vector(params).clone() xsms = [] t_start = timeit.default_timer() for k in range(self._max_k): with utils.Timer() as time: ## E-step xfilt, xfiltr, wfilt, meanll = self._fapf.filter(self._y) xsm = self._fapf.FFBSi(xfilt, wfilt) xsms.append(xsm) if self._xlen_cutoff: if len(xsms) > self._xlen_cutoff: xsms = xsms[-self._xlen_cutoff:] logger.logkv('train/Etime', time.dt) ## M-step with utils.Timer() as time: obj = lambda: -self.recursive_Q(xsms, self._y, 0, 0.) self._optimizer(obj, params) logger.logkv('train/Mtime', time.dt) logger.logkv('train/elapsedtime', timeit.default_timer() - t_start) ## log the current value of Q Q = float(self._fapf.Q_MCEM(self._y, xsms[-1])) logger.logkv('train/Q', Q) for callback in callbacks: callback(k) logger.dumpkvs() return params
def ecb(epoch): logger.logkv('time/epoch', epoch) params = list(system.parameters()) vparams = parameters_to_vector(params) error = (vparams - true_params).norm().item() logger.logkv('test/log10_paramerror', np.log10(error)) logger.logkv('time/epochtime', timeit.default_timer() - timer['start_time']) timer['start_time'] = timeit.default_timer() with torch.no_grad(): tgt_test_pr = system.step_derivs(t_test, x_test) error = float(torch.nn.functional.mse_loss(tgt_test_pr, tgt_test)) logger.logkv('test/log10_error', np.log10(error)) return
def train_net(net, train_data, train_trgt, test_data, test_trgt, valid_data, valid_trgt, y_std, lr, H, logdir, n_epochs=1000): T = train_data.shape[1] opt = torch.optim.Adam(net.parameters(), lr=lr) scheduler_off = 1000. scheduler = torch.optim.lr_scheduler.LambdaLR( opt, lambda epoch: scheduler_off / (scheduler_off + epoch)) best_val_loss = np.inf t0 = time.time() for e in range(n_epochs): logger.logkv('epoch', e) # train ll = [] coord_error = 0 for t in range(T - H + 1): opt.zero_grad() u = train_data[:, t:t + H] y = train_trgt[:, t + H - 1] y_pred = net(u) assert y.size() == y_pred.size() loss = compute_rms(y.unsqueeze(0), y_pred.unsqueeze(0), y_std) loss.backward() opt.step() ll.append(float(loss)) mean_train_loss = np.mean(ll) logger.logkv('log10_train_loss', np.log10(mean_train_loss)) coord_error /= (T - H) scheduler.step() for param_group in opt.param_groups: logger.logkv('log10_lr', np.log10(param_group['lr'])) if e % 100 == 0: # validation ll = [] coord_error = 0 for t in range(T - H): with torch.no_grad(): u = valid_data[:, t:t + H] y = valid_trgt[:, t + H - 1] y_pred = net(u) loss = compute_rms(y.unsqueeze(0), y_pred.unsqueeze(0), y_std) ll.append(float(loss)) mean_val_loss = np.mean(ll) logger.logkv('log10_val_loss', np.log10(mean_val_loss)) coord_error /= (T - H) # Test ll = [] coord_error = 0 for t in range(T - H): with torch.no_grad(): u = test_data[:, t:t + H] y = test_trgt[:, t + H - 1] y_pred = net(u) loss = compute_rms(y.unsqueeze(0), y_pred.unsqueeze(0), y_std) ll.append(float(loss)) mean_test_loss = np.mean(ll) logger.logkv('log10_test_loss', np.log10(mean_test_loss)) # Save if mean_val_loss < best_val_loss: torch.save(net.state_dict(), os.path.join(logdir, 'best_net.th')) best_val_loss = mean_val_loss if time.time() - t0 > 2: t0 = time.time() logger.dumpkvs() return net
def ecb(epoch): logger.logkv('test/rho', float(system._rho)) logger.logkv('test/sigma', float(system._sigma)) logger.logkv('test/beta', float(system._beta)) logger.logkv('test/rho_pcterr_log10', float(torch.log10((true_system._rho - system._rho).abs() / true_system._rho))) logger.logkv( 'test/sigma_pcterr_log10', float(torch.log10((true_system._sigma - system._sigma).abs() / true_system._sigma))) logger.logkv( 'test/beta_pcterr_log10', float(torch.log10((true_system._beta - system._beta).abs() / true_system._beta))) return
def ecb(epoch): logger.logkv('test/rho', float(system._rho)) logger.logkv('test/sigma', float(system._sigma)) logger.logkv('test/beta', float(system._beta)) logger.logkv('test/rho_pcterr_log10', float(torch.log10((true_system._rho - system._rho).abs() / true_system._rho))) logger.logkv( 'test/sigma_pcterr_log10', float(torch.log10((true_system._sigma - system._sigma).abs() / true_system._sigma))) logger.logkv( 'test/beta_pcterr_log10', float(torch.log10((true_system._beta - system._beta).abs() / true_system._beta))) logger.logkv('time/epochtime', timeit.default_timer() - timer['start_time']) timer['start_time'] = timeit.default_timer() return
def run(seed, lr, method, noise, damped, smoketest=False): """ Args: seed (int): seed lr (float): init learning rate method (str): training method in ['qdd', 'nqqd', 'del+mnorm', 'del+logdet'] noise (float): amount of noise in [0.01,0.05, 0.1] damped (bool): use damped pendulum data smoketest (bool): if smoketest, runs 2 epochs """ torch.set_default_dtype(torch.float64) dtype = torch.get_default_dtype() torch.manual_seed(seed) # load the data dataset = 'damped_' if damped else '' noisedict = { 0.01: '0p01', 0.05: '0p05', 0.10: '0p10', 0.20: '0p20', 0.30: '0p30', 0.40: '0p40', 0.50: '0p50', 1.0: '1p0' } dataset += 'dubpen_%s_smoothed.td' % noisedict[noise] dataset = './datasets/' + dataset data = torch.load('./datasets/%sdubpen_qddot.td' % ('damped_' if damped else '')) data_ = torch.load(dataset) dt = 0.05 logdir = 'data/%s_%s_%.1e_%.3f_%d' % ('damped' if damped else 'undamped', method, lr, noise, seed) logger.setup(logdir, action='d') inds = torch.randperm(16) Btr = 8 Bte = 4 Bva = 4 train_data_ = data_[inds[:Btr]] test_data_ = data_[inds[Btr:Btr + Bte]] val_data_ = data_[inds[Btr + Bte:Btr + Bte + Bva]] train_data = data[inds[:Btr]] test_data = data[inds[Btr:Btr + Bte]] val_data = data[inds[Btr + Bte:Btr + Bte + Bva]] t_, smq, smdq, smddq = train_data_[:] ttest_, smqtest, smdqtest, smddqtest = test_data_[:] tval_, smqval, smdqval, smddqval = val_data_[:] t, q, dq, ddq = train_data[:] ttest, qtest, dqtest, ddqtest = test_data[:] tval, qval, dqval, ddqval = val_data[:] B, T, qdim = q.shape # create the appropriate dataloader if 'del' in method: smq_1 = smq[:, :-2] smq_2 = smq[:, 1:-1] smq_3 = smq[:, 2:] smq_B = torch.stack([smq_1, smq_2, smq_3], dim=2).reshape(-1, 3, 2).detach() print(smq_B.shape, smq.shape) dataset = TensorDataset(smq_B) dataloader = DataLoader(dataset, batch_size=256, shuffle=True) elif method == 'qdd': dataset = TensorDataset(smq.reshape(-1, 1, 2), smdq.reshape(-1, 1, 2), smddq.reshape(-1, 1, 2)) dataloader = DataLoader(dataset, batch_size=256, shuffle=True) elif method == 'nqqd': x = torch.cat([smq, smdq], dim=-1) inp = x[:, :-1] out = x[:, 1:] inp = inp.reshape(-1, 1, 4) out = out.reshape(-1, 1, 4) dataset = TensorDataset(inp, out) dataloader = DataLoader(dataset, batch_size=256, shuffle=True) else: raise NotImplementedError # set up logdir and model if damped: system = ForcedSMM(qdim=qdim, dt=dt) else: system = StructuredMechanicalModel(qdim=qdim, dt=dt) # create the appropriate closure def qddcrit(system, smq_, smdq_, smddq_): ddq_ = system.compute_qddot(smq_, smdq_, create_graph=True) ddq_loss = torch.nn.functional.mse_loss(ddq_, smddq_) return ddq_loss def nqqdcrit(system, inp, out): out_ = system.step(torch.ones_like(inp)[..., 0], inp) nqqd_loss = torch.nn.functional.mse_loss(out_, out) return nqqd_loss if 'del' in method: dyncrit = DELCriterion(t_) if 'logdet' in method: bc = LogDetBarrierCriterion bcf = LogDetBarrierCriterion.mineig else: bc = MxNormBarrierCriterion bcf = MxNormBarrierCriterion.mmxnorm # initialize the barrier criterion, and find an appropriate coefficient between it and DELcrit lb = bcf(system, smq).detach() * 0.99 # interior point init barriercrit = bc(lb) delcrit = DELCriterion(t) with torch.no_grad(): dyncritloss = dyncrit(system, smq) barriercritloss = barriercrit(system, smq) mu = float(dyncritloss / barriercritloss) # mu makes them ~equal at init barriercrit = bc(lb, mu=mu, x_override=smq) crit = GroupCriterion([dyncrit, barriercrit]) elif method == 'qdd': crit = qddcrit elif method == 'nqqd': crit = nqqdcrit else: raise NotImplementedError # setup optimizer, scheduler opt = torch.optim.Adam(system.parameters(), lr=lr) sched = torch.optim.lr_scheduler.LambdaLR(opt, lambda k: 500. / (500 + k)) # train best_val_loss = np.inf best_val_loss_test_qddot = np.inf next_params = ptv(system.parameters()).detach() for epoch in range(2 if smoketest else 500): # train with SGD for batch in dataloader: prev_params = next_params opt.zero_grad() loss = crit(system, *batch) loss.backward() opt.step() if 'del' in method: # check line search n_ls = 0 while True: next_params = ptv(system.parameters()).detach() del_params = next_params - prev_params with torch.no_grad(): c = crit(system, smq) if torch.isnan(c): next_params = prev_params + 0.5 * del_params vtp(next_params, system.parameters()) n_ls += 1 else: break sched.step() with torch.no_grad(): val_sqmddqloss = qddcrit(system, smqval, smdqval, smddqval) train_qdd_loss = qddcrit(system, q, dq, ddq) test_qdd_loss = qddcrit(system, qtest, dqtest, ddqtest) val_qdd_loss = qddcrit(system, qval, dqval, ddqval) # select best model using validation error if val_sqmddqloss < best_val_loss: best_val_loss = float(val_sqmddqloss) best_val_loss_test_qddot_loss = float(test_qdd_loss) torch.save( system.state_dict(), os.path.join(logger.get_dir(), 'ckpts', 'best_model.th')) logger.logkv("train/epoch", epoch) logger.logkv("train/loss", float(loss)) logger.logkv("train/log10lr", np.log10(float(opt.param_groups[0]['lr']))) logger.logkv("eval/val_sqmddqloss", float(val_sqmddqloss)) logger.logkv("eval/train_qdd_loss", float(train_qdd_loss)) logger.logkv("eval/test_qdd_loss", float(test_qdd_loss)) logger.logkv("eval/val_qdd_loss", float(val_qdd_loss)) logger.logkv("eval/best_val_loss", float(best_val_loss)) logger.logkv("eval/best_val_loss_test_qddot_loss", float(best_val_loss_test_qddot_loss)) logger.dumpkvs()