def train_toy(toy, load=True, nb_steps=20, nb_flow=1, folder=""): device = "cpu" logger = utils.get_logger(logpath=os.path.join(folder, toy, 'logs'), filepath=os.path.abspath(__file__)) logger.info("Creating model...") model = UMNNMAFFlow(nb_flow=nb_flow, nb_in=2, hidden_derivative=[50, 50, 50, 50], hidden_embedding=[50, 50, 50, 50], embedding_s=10, nb_steps=nb_steps, device=device).to(device) logger.info("Model created.") opt = torch.optim.Adam(model.parameters(), 1e-3, weight_decay=1e-5) if load: logger.info("Loading model...") model.load_state_dict(torch.load(folder + toy + '/model.pt')) model.train() opt.load_state_dict(torch.load(folder + toy + '/ADAM.pt')) logger.info("Model loaded.") nb_samp = 1000 batch_size = 100 x_test = torch.tensor(toy_data.inf_train_gen(toy, batch_size=1000)).to(device) x = torch.tensor(toy_data.inf_train_gen(toy, batch_size=1000)).to(device) for epoch in range(10000): ll_tot = 0 start = timer() for j in range(0, nb_samp, batch_size): cur_x = torch.tensor( toy_data.inf_train_gen(toy, batch_size=batch_size)).to(device) ll, z = model.compute_ll(cur_x) ll = -ll.mean() ll_tot += ll.detach() / (nb_samp / batch_size) loss = ll opt.zero_grad() loss.backward() opt.step() end = timer() ll_test, _ = model.compute_ll(x_test) ll_test = -ll_test.mean() logger.info( "epoch: {:d} - Train loss: {:4f} - Test loss: {:4f} - Elapsed time per epoch {:4f} (seconds)" .format(epoch, ll_tot.item(), ll_test.item(), end - start)) if (epoch % 100) == 0: summary_plots(x, x_test, folder, epoch, model, ll_tot, ll_test) torch.save(model.state_dict(), folder + toy + '/model.pt') torch.save(opt.state_dict(), folder + toy + '/ADAM.pt')
def compute_loss(args, model, batch_size=None, beta=1.): if batch_size is None: batch_size = args.batch_size # load data x = toy_data.inf_train_gen(args.data, batch_size=batch_size) x = torch.from_numpy(x).type(torch.float32).to(device) #print(x.shape) #print(x) #(500, 2) #tensor([[3.4526e+00, 1.4150e-01], # [3.1749e+00, 7.9239e-01], # [3.4892e+00, -3.4557e-01], # (...) # [2.9311e+00, 1.4599e-01], # [2.7587e+00, 3.3307e-02], # [2.4231e+00, -1.2663e-01]], device = 'cuda:0') zero = torch.zeros(x.shape[0], 1).to(x) # transform to z z, delta_logp = model(x, zero) # compute log p(z) logpz = standard_normal_logprob(z).sum(1, keepdim=True) logpx = logpz - beta * delta_logp loss = -torch.mean(logpx) return loss, torch.mean(logpz), torch.mean(-delta_logp)
def loss_fn2(first_term_loss, genFGen2, args, model): # diff = t1 - t2 # .numel() number of elements # return torch.sum(diff * diff) / diff.numel() xData = toy_data.inf_train_gen(args.data, batch_size=args.batch_size) xData = torch.from_numpy(xData).type(torch.float32).to(device) # second_term_loss = 0.0 # second_term_loss = torch.from_numpy(np.array(0.0, dtype='float32')) # second_term_loss = torch.from_numpy(np.array(0.0, dtype='float32')) second_term_loss = torch.from_numpy(np.array(0.0, dtype='float32')) for i in genFGen2: for j in xData: # second_term_loss += np.linalg.norm(i.cpu().detach().numpy()-j.cpu().detach().numpy()) # second_term_loss += torch.norm(i-j, 2) # second_term_loss += torch.norm(i-j) # second_term_loss += torch.dist(i.type(torch.float64), j.type(torch.float64), 2) second_term_loss += torch.dist(i, j, 2) second_term_loss /= args.batch_size second_term_loss *= 30 second_term_loss /= args.batch_size # second_term_loss = torch.tensor(second_term_loss).to(device) #return first_term_loss + second_term_loss + third_term_loss #return first_term_loss # third_term_loss = 0.0 # third_term_loss = torch.from_numpy(np.array(0.0, dtype='float32')) # third_term_loss = torch.from_numpy(np.array(0.0, dtype='float32')) third_term_loss = torch.from_numpy(np.array(0.0, dtype='float32')) for i in range(args.batch_size): for j in range(args.batch_size): if i != j: # third_term_loss += ((np.linalg.norm(genFGen3[i,:].cpu().detach().numpy()-genFGen3[j,:].cpu().detach().numpy())) / (np.linalg.norm(genFGen2[i,:].cpu().detach().numpy()-genFGen2[j,:].cpu().detach().numpy()))) # third_term_loss += ((torch.norm(genFGen3[i,:]-genFGen3[j,:], 2)) / (torch.norm(genFGen2[i,:]-genFGen2[j,:], 2))) # third_term_loss += ((torch.norm(genFGen3[i,:]-genFGen3[j,:])) / (torch.norm(genFGen2[i,:]-genFGen2[j,:]))) # third_term_loss += ((torch.norm(genFGen3[i,:] - genFGen3[j,:])) / (torch.norm(genFGen2[i,:] - genFGen2[j,:]))) third_term_loss += ( (torch.dist(genFGen3[i, :], genFGen3[j, :], 2)) / (torch.dist(genFGen2[i, :], genFGen2[j, :], 2))) third_term_loss /= (args.batch_size - 1) third_term_loss *= 10 third_term_loss /= args.batch_size # third_term_loss = torch.tensor(third_term_loss).to(device) return first_term_loss + second_term_loss + third_term_loss
def get_ckpt_model_and_data(args): # Load checkpoint. checkpt = torch.load(args.checkpt, map_location=lambda storage, loc: storage) ckpt_args = checkpt['args'] state_dict = checkpt['state_dict'] # Construct model and restore checkpoint. regularization_fns, regularization_coeffs = create_regularization_fns( ckpt_args) model = build_model_tabular(ckpt_args, 2, regularization_fns).to(device) if ckpt_args.spectral_norm: add_spectral_norm(model) set_cnf_options(ckpt_args, model) model.load_state_dict(state_dict) model.to(device) print(model) print("Number of trainable parameters: {}".format( count_parameters(model))) # Load samples from dataset data_samples = toy_data.inf_train_gen(ckpt_args.data, batch_size=2000) return model, data_samples
def compute_loss(args, model, batch_size=args.batch_size): x = toy_data.inf_train_gen(args.data, batch_size=batch_size) x = torch.from_numpy(x).type(torch.float32).to(device) zero = torch.zeros(x.shape[0], 1).to(x) z, change = model(x, zero) logpx = standard_normal_logprob(z).sum(1, keepdim=True) - change loss = -torch.mean(logpx) return loss
def compute_loss(args, model, batch_size=None, beta=1.): if batch_size is None: batch_size = args.batch_size # load data x = toy_data.inf_train_gen(args.data, batch_size=batch_size) x = torch.from_numpy(x).type(torch.float32).to(device) zero = torch.zeros(x.shape[0], 1).to(x) # transform to z z, delta_logp = model(x, zero) # compute log p(z) logpz = standard_normal_logprob(z).sum(1, keepdim=True) logpx = logpz - beta * delta_logp loss = -torch.mean(logpx) return loss, torch.mean(logpz), torch.mean(-delta_logp)
def compute_loss(args, model, batch_size=None): if batch_size is None: batch_size = args.batch_size # load data x = toy_data.inf_train_gen(args.data, batch_size=batch_size) x = torch.from_numpy(x).type(torch.float32).to(device) zero = torch.zeros(x.shape[0], 1).to(x) lec = None if (args.poly_coef is None or not model.training) else torch.tensor(0.0).to(x) # transform to z z, delta_logp, lec = model(x, zero, lec) # compute log q(z) logpz = standard_normal_logprob(z).sum(1, keepdim=True) logpx = logpz - delta_logp loss = -torch.mean(logpx) return loss, lec
def compute_loss(args, model, batch_size=None): if batch_size is None: batch_size = args.batch_size # load data x = toy_data.inf_train_gen(args.data, batch_size=batch_size) x = torch.from_numpy(x).type(torch.float32).to(device) zero = torch.zeros(x.shape[0], 1).to(x) # transform to z std = (args.std_max - args.std_min) * torch.rand_like(x[:,0]).view(-1,1) + args.std_min eps = torch.randn_like(x) * std std_in = std / args.std_max * args.std_weight z, delta_logp = model(x+eps, std_in, zero) # compute log q(z) logpz = standard_normal_logprob(z).sum(1, keepdim=True) logpx = logpz - delta_logp loss = -torch.mean(logpx) return loss
def compute_loss(args, model, batch_size=None, beta=1.): if batch_size is None: batch_size = args.batch_size # load data # load data x = toy_data.inf_train_gen(args.data, batch_size=batch_size) x = torch.from_numpy(x).type(torch.float32).to(device) #print('') #print(x.shape) #print('') #print(x) #print('') #plt.figure() #plt.plot(x[:, 0].cpu().squeeze().numpy(), x[:, 1].cpu().squeeze().numpy(), 'o') #plt.ion() #plt.show() #plt.pause(2) #sfadadfa zero = torch.zeros(x.shape[0], 1).to(x) # transform to z z, delta_logp = model(x, zero) # compute log p(z) # compute log p(z) logpz = standard_normal_logprob(z).sum(1, keepdim=True) logpx = logpz - beta * delta_logp loss = -torch.mean(logpx) return loss, torch.mean(logpz), torch.mean(-delta_logp)
def loss_fn2(genFGen2, args, model): genFGen3 = torch.randn((args.batch_size, 2)).to(device) first_term_loss = compute_loss2(genFGen2, args, model, beta=beta) xData = toy_data.inf_train_gen(args.data, batch_size=args.batch_size) xData = torch.from_numpy(xData).type(torch.float32).to(device) #second_term_loss = 0.0 #second_term_loss = torch.from_numpy(np.array(0.0, dtype='float32')) #second_term_loss = torch.from_numpy(np.array(0.0, dtype='float32')) #second_term_loss = torch.from_numpy(np.array(0.0, dtype='float32')) #second_term_loss = torch.from_numpy(np.array(0.0, dtype='float32')) #second_term_loss = torch.from_numpy(np.array(0.0, dtype='float32')).to(device) #second_term_loss = torch.from_numpy(np.array(0.0, dtype='float32')).to(device) #second_term_loss = torch.from_numpy(np.array(0.0, dtype='float32')).to(device) second_term_loss = torch.from_numpy(np.array(float('inf'), dtype='float32')).to(device) for i in genFGen2: for j in xData: # second_term_loss += np.linalg.norm(i.cpu().detach().numpy()-j.cpu().detach().numpy()) # second_term_loss += torch.norm(i-j, 2) # second_term_loss += torch.norm(i-j) # second_term_loss += torch.dist(i.type(torch.float64), j.type(torch.float64), 2) #second_term_loss += torch.dist(i, j, 2) store_second_term_loss = torch.dist(i, j, 2) if second_term_loss < store_second_term_loss: second_term_loss = store_second_term_loss #second_term_loss /= args.batch_size second_term_loss *= 0.1 second_term_loss /= args.batch_size # second_term_loss = torch.tensor(second_term_loss).to(device) #return first_term_loss + second_term_loss + third_term_loss #return first_term_loss #third_term_loss = 0.0 #third_term_loss = torch.from_numpy(np.array(0.0, dtype='float32')) #third_term_loss = torch.from_numpy(np.array(0.0, dtype='float32')) #third_term_loss = torch.from_numpy(np.array(0.0, dtype='float32')) #third_term_loss = torch.from_numpy(np.array(0.0, dtype='float32')) third_term_loss = torch.from_numpy(np.array(0.0, dtype='float32')).to(device) for i in range(args.batch_size): for j in range(args.batch_size): if i != j: # third_term_loss += ((np.linalg.norm(genFGen3[i,:].cpu().detach().numpy()-genFGen3[j,:].cpu().detach().numpy())) / (np.linalg.norm(genFGen2[i,:].cpu().detach().numpy()-genFGen2[j,:].cpu().detach().numpy()))) # third_term_loss += ((torch.norm(genFGen3[i,:]-genFGen3[j,:], 2)) / (torch.norm(genFGen2[i,:]-genFGen2[j,:], 2))) # third_term_loss += ((torch.norm(genFGen3[i,:]-genFGen3[j,:])) / (torch.norm(genFGen2[i,:]-genFGen2[j,:]))) # third_term_loss += ((torch.norm(genFGen3[i,:] - genFGen3[j,:])) / (torch.norm(genFGen2[i,:] - genFGen2[j,:]))) third_term_loss += ( (torch.dist(genFGen3[i, :], genFGen3[j, :], 2)) / (torch.dist(genFGen2[i, :], genFGen2[j, :], 2))) third_term_loss /= (args.batch_size - 1) third_term_loss *= 0.1 third_term_loss /= args.batch_size # third_term_loss = torch.tensor(third_term_loss).to(device) return first_term_loss + second_term_loss + third_term_loss
ax.set_facecolor(COLOR_BACK) fig_filename = os.path.join(save_path, file_name + '.png') utils.makedirs(os.path.dirname(fig_filename)) plt.savefig(fig_filename, format='png', dpi=1200, bbox_inches='tight') plt.close() if __name__ == '__main__': save_path = 'generate1/' + args.data n_samples = 2000 COLOR = 1, 1, 1, 0.64 if not os.path.exists(save_path): os.makedirs(save_path) sample_real = toy_data.inf_train_gen(args.data, batch_size=n_samples) save_fig(sample_real, COLOR, 'D', 0.1, save_path, 'sample_data') softflow = build_model_tabular(args, 2).to(device) softflow_path = args.load_path ckpt_softflow = torch.load(softflow_path) softflow.load_state_dict(ckpt_softflow['state_dict']) softflow.eval() z = torch.randn(n_samples, 2).type(torch.float32).to(device) sample_s = [] inds = torch.arange(0, z.shape[0]).to(torch.int64) with torch.no_grad(): for ii in torch.split(inds, int(100**2)): zeros_std = torch.zeros(z[ii].shape[0], 1).to(z) sample_s.append(softflow(z[ii], zeros_std, reverse=True))
def sample(self): return torch.tensor(inf_train_gen(self.data, batch_size=1)), torch.tensor([])
def loss_fn2(genFGen2, args, model): genFGen3 = torch.randn((args.batch_size, 2)).to(device) #first_term_loss = compute_loss2(genFGen2, args, model, beta=beta) #first_term_loss = compute_loss2(genFGen2, args, model, beta=beta) #first_term_loss = compute_loss2(genFGen2, args, model, beta=beta) #print('') #print(first_term_loss) import math #mu = torch.from_numpy(np.array([2.8099582e+00, 9.6440443e-04], dtype="float32")).to(device) #mu = torch.from_numpy(np.array([[2.805741, -0.00889241]], dtype="float32")).to(device) mu = torch.from_numpy(np.array([2.805741, -0.00889241], dtype="float32")).to(device) #S = torch.from_numpy(np.array([[0.35833913, 0.0], [0.0, 0.34720358]], dtype="float32")).to(device) #S = torch.from_numpy(np.array([[pow(0.35833913,2), 0.0], [0.0, pow(0.34720358,2)]], dtype="float32")).to(device) S = torch.from_numpy( np.array([[pow(0.3442525, 2), 0.0], [0.0, pow(0.35358343, 2)]], dtype="float32")).to(device) #print('') #import timeit #import scipy.stats #start = timeit.default_timer() #storeAll = torch.from_numpy(np.array(0.0, dtype="float32")).to(device) #for loopIndex_i in range(genFGen2.size()[0]): # #print((((torch.from_numpy(np.array(-np.log(2 * math.pi), dtype="float32")).to(device))) - # # ((0.5 * torch.log(torch.det(S)))) - 0.5 * # # ((((torch.matmul(torch.matmul(((genFGen2[loopIndex_i:1+loopIndex_i, :]) - mu), torch.inverse(S)), # # torch.transpose((genFGen2[loopIndex_i:1+loopIndex_i, :]) - mu, 0, 1)))))))) # # storeAll += ((((torch.from_numpy(np.array(-np.log(2 * math.pi), dtype="float32")).to(device))) - # ((0.5 * torch.log(torch.det(S)))) - 0.5 * ( # (((torch.matmul(torch.matmul(((genFGen2[loopIndex_i:1 + loopIndex_i, :]) - mu), torch.inverse(S)), # torch.transpose((genFGen2[loopIndex_i:1 + loopIndex_i, :]) - mu, 0, 1)))))))).squeeze() # #storeAll /= genFGen2.size()[0] #print(storeAll) #stop = timeit.default_timer() #print('Time: ', stop - start) #print('') #print(torch.exp(((torch.from_numpy(np.array(-np.log(2 * math.pi), dtype="float32")).to(device))) - # ((0.5 * torch.log(torch.det(S)))) - 0.5 * # ((((torch.matmul(torch.matmul(((genFGen2[0:1, :]) - mu), torch.inverse(S)), # torch.transpose((genFGen2[0:1, :]) - mu, 0, 1)))))))) #import scipy.stats #print(scipy.stats.multivariate_normal.pdf(genFGen2[0:1, :].cpu().detach().numpy(), # mean=mu.cpu().detach().numpy(), cov=S.cpu().detach().numpy())) #first_term_loss = storeAll #asdfasdfadfa #print('') #import timeit #import scipy.stats #start = timeit.default_timer() storeAll = torch.from_numpy(np.array(0.0, dtype="float32")).to(device) for loopIndex_i in range(genFGen2.size()[0]): #storeAll += np.log(scipy.stats.multivariate_normal.pdf(genFGen2[loopIndex_i:1+loopIndex_i, :].cpu().detach().numpy(), # mean=mu.cpu().detach().numpy(), cov=S.cpu().detach().numpy())) toUse_storeAll = torch.distributions.MultivariateNormal( loc=mu, covariance_matrix=S) storeAll += toUse_storeAll.log_prob( genFGen2[loopIndex_i:1 + loopIndex_i, :].squeeze(0)) storeAll /= genFGen2.size()[0] #first_term_loss = storeAll #print(storeAll) #asdfasdfadfa #first_term_loss = storeAll #first_term_loss = storeAll #first_term_loss = 1.0e16 * torch.exp(storeAll) #first_term_loss = storeAll first_term_loss = storeAll #stop = timeit.default_timer() #print('Time: ', stop - start) #print('') #asdfsdfaa xData = toy_data.inf_train_gen(args.data, batch_size=args.batch_size) xData = torch.from_numpy(xData).type(torch.float32).to(device) #import timeit #start = timeit.default_timer() #second_term_loss = 0.0 #second_term_loss = torch.from_numpy(np.array(0.0, dtype='float32')) #second_term_loss = torch.from_numpy(np.array(0.0, dtype='float32')) #second_term_loss = torch.from_numpy(np.array(0.0, dtype='float32')) #second_term_loss = torch.from_numpy(np.array(0.0, dtype='float32')) #second_term_loss = torch.from_numpy(np.array(0.0, dtype='float32')).to(device) #second_term_loss = torch.from_numpy(np.array(0.0, dtype='float32')).to(device) #for i in genFGen2: # for j in xData: # # second_term_loss += np.linalg.norm(i.cpu().detach().numpy()-j.cpu().detach().numpy()) # # # second_term_loss += torch.norm(i-j, 2) # # second_term_loss += torch.norm(i-j) # # # second_term_loss += torch.dist(i.type(torch.float64), j.type(torch.float64), 2) # second_term_loss += torch.dist(i, j, 2) # # second_term_loss /= args.batch_size # #second_term_loss *= 0.1 # #second_term_loss /= args.batch_size #print('') #print(first_term_loss) #print(torch.exp(first_term_loss)) #print('') #print(second_term_loss) #stop = timeit.default_timer() #print('Time: ', stop - start) #print('') #asdfsdfa #second_term_loss = [] #for i in genFGen2: # second_term_lossTwo = [] # for j in xData: # #second_term_lossTwo.append(torch.dist(i, j, 2).unsqueeze(0)) # second_term_lossTwo.append(torch.dist(i, j, 2)) # second_term_loss.append(second_term_lossTwo) # #second_term_loss.append(.unsqueeze(0)) # #second_term_loss.append() #second_term_loss = torch.cat(second_term_loss) ##second_term_loss = torch.cat(second_term_loss, dim=1) ##second_term_loss *= 0.1 #second_term_loss = torch.mean(second_term_loss) #import timeit #start = timeit.default_timer() var1 = [] for i in genFGen2: for j in xData: new_stuff = torch.dist(i, j, 2) # this is a tensor var1.append(new_stuff.unsqueeze(0)) #var1.append(new_stuff) var1_tensor = torch.cat(var1) #var1_tensor = torch.cat(var1, dim=1) #second_term_loss = torch.mean(var1_tensor) #second_term_loss = torch.mean(var1_tensor) / args.batch_size #second_term_loss = 10.0 * torch.mean(var1_tensor) / args.batch_size #second_term_loss = torch.min(var1_tensor) second_term_loss = torch.min(var1_tensor) / args.batch_size #second_term_loss = 10.0 * torch.min(var1_tensor) / args.batch_size #second_term_loss = torch.from_numpy(np.array(None, dtype='float32')).to(device) #for i in genFGen2: # second_term_lossTwo = torch.from_numpy(np.array(None, dtype='float32')).to(device) # for j in xData: # second_term_lossTwo = torch.cat((second_term_lossTwo, torch.dist(i, j, 2)), 0) # second_term_loss = torch.cat((second_term_loss, torch.mean(second_term_lossTwo, 0)), 0) # #second_term_loss *= 0.1 #second_term_loss = torch.mean(second_term_loss) #print('') #print(second_term_loss) #stop = timeit.default_timer() #print('Time: ', stop - start) #print('') #print(second_term_loss) #print(torch.min(var1_tensor) / args.batch_size) #print('') #asdfdffa second_term_loss *= 10000.0 #second_term_loss *= 1.0e-13 #second_term_loss *= 1.0e3 if torch.abs(first_term_loss) / second_term_loss > 200.0: second_term_loss /= 2.0 print('') print(first_term_loss) print(second_term_loss) print('') #print('') #print(torch.exp(first_term_loss)) #print(1.0e-13 * second_term_loss/10000.0) #print('') #asdfdffa #third_term_loss = torch.from_numpy(np.array(0.0, dtype='float32')).to(device) #for i in range(args.batch_size): # for j in range(args.batch_size): # if i != j: # # third_term_loss += ((np.linalg.norm(genFGen3[i,:].cpu().detach().numpy()-genFGen3[j,:].cpu().detach().numpy())) / (np.linalg.norm(genFGen2[i,:].cpu().detach().numpy()-genFGen2[j,:].cpu().detach().numpy()))) # # # third_term_loss += ((torch.norm(genFGen3[i,:]-genFGen3[j,:], 2)) / (torch.norm(genFGen2[i,:]-genFGen2[j,:], 2))) # # third_term_loss += ((torch.norm(genFGen3[i,:]-genFGen3[j,:])) / (torch.norm(genFGen2[i,:]-genFGen2[j,:]))) # # # third_term_loss += ((torch.norm(genFGen3[i,:] - genFGen3[j,:])) / (torch.norm(genFGen2[i,:] - genFGen2[j,:]))) # third_term_loss += ((torch.dist(genFGen3[i, :], genFGen3[j, :], 2)) / (torch.dist(genFGen2[i, :], genFGen2[j, :], 2))) # third_term_loss /= (args.batch_size - 1) #third_term_loss /= args.batch_size #third_term_loss *= 100.0 #print(third_term_loss) #print('') #asdfsfaa #return first_term_loss + second_term_loss + third_term_loss return first_term_loss + second_term_loss
if test_loss.item() < best_loss: best_loss = test_loss.item() utils.makedirs(args.save) torch.save({ 'args': args, 'state_dict': model.state_dict(), }, os.path.join(args.save, 'checkpt.pth')) model.train() #if itr == 1 or itr % args.viz_freq == 0: if itr % args.viz_freq == 0: with torch.no_grad(): model.eval() p_samples = toy_data.inf_train_gen(args.data, batch_size=20000) sample_fn, density_fn = model.inverse, model.forward plt.figure(figsize=(9, 3)) visualize_transform(p_samples, torch.randn, standard_normal_logprob, transform=sample_fn, inverse_transform=density_fn, samples=True, npts=400, device=device) fig_filename = os.path.join(args.save, 'figs', '{:04d}.jpg'.format(itr)) print('') print(fig_filename) print('') utils.makedirs(os.path.dirname(fig_filename)) plt.savefig(fig_filename)
log_message = '[TEST] Iter {:04d} | Test Loss {:.6f} | NFE {:.0f}'.format(itr, test_loss, test_nfe) logger.info(log_message) if test_loss.item() < best_loss: best_loss = test_loss.item() utils.makedirs(args.save) torch.save({ 'args': args, 'state_dict': model.state_dict(), }, os.path.join(args.save, 'checkpt.pth')) model.train() if itr % args.viz_freq == 0: with torch.no_grad(): model.eval() p_samples = toy_data.inf_train_gen(args.data, rng_seed=420000, batch_size=2000) sample_fn, density_fn = get_transforms(model) plt.figure(figsize=(9, 3)) visualize_transform( p_samples, torch.randn, standard_normal_logprob, transform=sample_fn, inverse_transform=density_fn, samples=True, npts=800, device=device ) fig_filename = os.path.join(args.save, 'figs', '{:04d}.pdf'.format(itr)) utils.makedirs(os.path.dirname(fig_filename)) plt.savefig(fig_filename) plt.close() model.train() end = time.time()
def loss_fn2(genFGen2, args, model): first_term_loss = compute_loss2(genFGen2, args, model) #first_term_loss2 = compute_loss2(genFGen2, args, model) #first_term_loss = torch.log(first_term_loss2 / (1.0 - first_term_loss2)) #print('') #print(first_term_loss) #mu = torch.from_numpy(np.array([2.805741, -0.00889241], dtype="float32")).to(device) #S = torch.from_numpy(np.array([[pow(0.3442525,2), 0.0], [0.0, pow(0.35358343,2)]], dtype="float32")).to(device) #storeAll = torch.from_numpy(np.array(0.0, dtype="float32")).to(device) #toUse_storeAll = torch.distributions.MultivariateNormal(loc=mu, covariance_matrix=S) #for loopIndex_i in range(genFGen2.size()[0]): # storeAll += torch.exp(toUse_storeAll.log_prob(genFGen2[loopIndex_i:1 + loopIndex_i, :].squeeze(0))) #storeAll /= genFGen2.size()[0] #print(storeAll) #print('') #print('') #print(compute_loss2(mu.unsqueeze(0), args, model)) #print(torch.exp(toUse_storeAll.log_prob(mu))) #print('') #first_term_loss = storeAll xData = toy_data.inf_train_gen(args.data, batch_size=args.batch_size) xData = torch.from_numpy(xData).type(torch.float32).to(device) #var2 = [] #for i in genFGen2: # var1 = [] # for j in xData: # new_stuff = torch.dist(i, j, 2) # this is a tensor # var1.append(new_stuff.unsqueeze(0)) # var1_tensor = torch.cat(var1) # second_term_loss2 = torch.min(var1_tensor) / args.batch_size # var2.append(second_term_loss2.unsqueeze(0)) #var2_tensor = torch.cat(var2) #second_term_loss = torch.mean(var2_tensor) / args.batch_size #second_term_loss *= 100.0 #print('') #print(second_term_loss) # If you know in advance the size of the final tensor, you can allocate # an empty tensor beforehand and fill it in the for loop. #x = torch.empty(size=(len(items), 768)) #for i in range(len(items)): # x[i] = calc_result #print(len(genFGen2)) #print(genFGen2.shape[0]) # len(.) and not .shape[0] #print(len(xData)) #print(xData.shape[0]) # Use len(.) and not .shape[0] #second_term_loss = torch.empty(size=(len(genFGen2), len(xData))).to(device) #second_term_loss = torch.empty(size=(len(genFGen2), len(xData)), device=device, requires_grad=True) #second_term_loss3 = torch.empty(size=(len(genFGen2), len(xData)), device=device, requires_grad=True) second_term_loss3 = torch.empty(size=(len(genFGen2), len(xData)), device=device, requires_grad=False) for i in range(len(genFGen2)): for j in range(len(xData)): #second_term_loss[i, j] = torch.dist(genFGen2[i,:], xData[j,:], 2) #second_term_loss[i, j] = torch.dist(genFGen2[i, :], xData[j, :], 1) #second_term_loss3[i, j] = torch.dist(genFGen2[i, :], xData[j, :], 1) #second_term_loss3[i, j] = torch.dist(genFGen2[i, :], xData[j, :], 1) #second_term_loss3[i, j] = torch.dist(genFGen2[i, :], xData[j, :], 1) #second_term_loss3[i, j] = torch.dist(genFGen2[i, :], xData[j, :], 1) #second_term_loss3[i, j] = torch.tensor(0.1, requires_grad=True) #second_term_loss3[i, j] = torch.dist(genFGen2[i, :], xData[j, :], 1) #second_term_loss3[i, j] = torch.dist(genFGen2[i, :], xData[j, :], 1).requires_grad_() #second_term_loss3[i, j] = torch.dist(genFGen2[i, :], xData[j, :], 1).requires_grad_() second_term_loss3[i, j] = torch.dist(genFGen2[i, :], xData[j, :], 2).requires_grad_()**2 #second_term_loss[i, j] = torch.dist(genFGen2[i, :], xData[j, :], 2)**2 #second_term_loss2, _ = torch.min(second_term_loss, 1) second_term_loss2, _ = torch.min(second_term_loss3, 1) second_term_loss = 1000.0 * torch.mean(second_term_loss2) / ( args.batch_size**2) #print(second_term_loss) #print('') print('') print(first_term_loss) print(second_term_loss) #third_term_loss = torch.from_numpy(np.array(0.0, dtype='float32')).to(device) #for i in range(args.batch_size): # for j in range(args.batch_size): # if i != j: # # third_term_loss += ((np.linalg.norm(genFGen3[i,:].cpu().detach().numpy()-genFGen3[j,:].cpu().detach().numpy())) / (np.linalg.norm(genFGen2[i,:].cpu().detach().numpy()-genFGen2[j,:].cpu().detach().numpy()))) # # # third_term_loss += ((torch.norm(genFGen3[i,:]-genFGen3[j,:], 2)) / (torch.norm(genFGen2[i,:]-genFGen2[j,:], 2))) # # third_term_loss += ((torch.norm(genFGen3[i,:]-genFGen3[j,:])) / (torch.norm(genFGen2[i,:]-genFGen2[j,:]))) # # # third_term_loss += ((torch.norm(genFGen3[i,:] - genFGen3[j,:])) / (torch.norm(genFGen2[i,:] - genFGen2[j,:]))) # third_term_loss += ((torch.dist(genFGen3[i, :], genFGen3[j, :], 2)) / (torch.dist(genFGen2[i, :], genFGen2[j, :], 2))) # third_term_loss /= (args.batch_size - 1) #third_term_loss /= args.batch_size ##third_term_loss *= 1000.0 #print(third_term_loss) #print('') print('') #asdfsfa #return first_term_loss + second_term_loss + third_term_loss #return first_term_loss + second_term_loss #return second_term_loss return first_term_loss + second_term_loss
x = torch.from_numpy(x).type(torch.float32).to(device) zero = torch.zeros(x.shape[0], 1).to(x) # transform to z z, delta_logp = model(x, zero) # compute log q(z) logpz = standard_normal_logprob(z).sum(1, keepdim=True) logpx = logpz - delta_logp loss = -torch.mean(logpx) return loss if __name__ == '__main__': x = toy_data.inf_train_gen(args.data, batch_size=args.batch_size) print(x.shape) plt.hist(x, bins=500) plt.savefig('testwill.png') regularization_fns, regularization_coeffs = create_regularization_fns(args) # model = build_model_tabular(args, 200, regularization_fns).to(device) model = build_model_tabular(args, 1, regularization_fns).to(device) if args.spectral_norm: add_spectral_norm(model) set_cnf_options(args, model) # logger.info(model) logger.info("Number of trainable parameters: {}".format( count_parameters(model))) optimizer = optim.Adam(model.parameters(),
alph = checkpt['args'].alph nTh = checkpt['args'].nTh d = checkpt['state_dict']['A'].size(1) - 1 net = Phi(nTh=nTh, m=m, d=d, alph=alph) prec = checkpt['state_dict']['A'].dtype net = net.to(prec) net.load_state_dict(checkpt['state_dict']) net = net.to(device) args.data = checkpt['args'].data torch.set_default_dtype(prec) cvt = lambda x: x.type(prec).to(device, non_blocking=True) nSamples = args.batch_size p_samples = cvt(torch.Tensor(toy_data.inf_train_gen(args.data, batch_size=nSamples))) y = cvt(torch.randn(nSamples,d)) net.eval() with torch.no_grad(): test_loss, test_cs = compute_loss(net, p_samples, args.nt) # sample_fn, density_fn = get_transforms(model) modelFx = integrate(p_samples[:, 0:d], net, [0.0, 1.0], args.nt, stepper="rk4", alph=net.alph) modelFinvfx = integrate(modelFx[:, 0:d] , net, [1.0, 0.0], args.nt, stepper="rk4", alph=net.alph) modelGen = integrate(y[:, 0:d] , net, [1.0, 0.0], args.nt, stepper="rk4", alph=net.alph) print(" {:9s} {:9s} {:11s} {:9s}".format( "loss", "L (L_2)", "C (loss)", "R (HJB)")) print("[TEST]: {:9.3e} {:9.3e} {:11.5e} {:9.3e}".format(test_loss, test_cs[0], test_cs[1], test_cs[2]))
def summary_plots(x, x_test, folder, epoch, model, ll_tot, ll_test): fig = plt.figure() ax = plt.subplot(1, 3, 1, aspect="equal") vf.plt_flow(model.compute_ll, ax) ax = plt.subplot(1, 3, 2, aspect="equal") vf.plt_samples(toy_data.inf_train_gen(toy, batch_size=1000000), ax, npts=200) ax = plt.subplot(1, 3, 3, aspect="equal") samples = model.invert( torch.distributions.Normal(0., 1.).sample([1000, 2]), 8, "Binary") vf.plt_samples(samples.detach().numpy(), ax, title="$x\sim q(x)$") plt.savefig("%s/flow_%d.png" % (folder + toy, epoch)) plt.close(fig) fig = plt.figure() z = torch.distributions.Normal(0., 1.).sample(x_test.shape) plt.subplot(222) plt.title("z") plt.scatter(z[:, 0], z[:, 1], alpha=.2) x_min = z.min(0)[0] - .5 x_max = z.max(0)[0] + .5 ticks = [1, 1] plt.xticks(np.arange(int(x_min[0]), int(x_max[0]), ticks[0]), np.arange(int(x_min[0]), int(x_max[0]), ticks[0])) plt.yticks(np.arange(int(x_min[1]), int(x_max[1]), ticks[1]), np.arange(int(x_min[1]), int(x_max[1]), ticks[1])) z_pred = model.forward(x_test) z_pred = z_pred.detach().cpu().numpy() plt.subplot(221) plt.title("z pred") plt.scatter(z_pred[:, 0], z_pred[:, 1], alpha=.2) plt.xticks(np.arange(int(x_min[0]), int(x_max[0]), ticks[0]), np.arange(int(x_min[0]), int(x_max[0]), ticks[0])) plt.yticks(np.arange(int(x_min[1]), int(x_max[1]), ticks[1]), np.arange(int(x_min[1]), int(x_max[1]), ticks[1])) start = timer() z = torch.distributions.Normal(0., 1.).sample((1000, 2)) x_pred = model.invert(z, 5, "ParallelSimpler") end = timer() print("Inversion time: {:4f}s".format(end - start)) plt.subplot(223) plt.title("x pred") x_pred = x_pred.detach().cpu().numpy() plt.scatter(x_pred[:, 0], x_pred[:, 1], alpha=.2) x_min = x.min(0)[0] - .5 x_max = x.max(0)[0] + .5 ticks = [1, 1] plt.xticks(np.arange(int(x_min[0]), int(x_max[0]), ticks[0]), np.arange(int(x_min[0]), int(x_max[0]), ticks[0])) plt.yticks(np.arange(int(x_min[1]), int(x_max[1]), ticks[1]), np.arange(int(x_min[1]), int(x_max[1]), ticks[1])) plt.subplot(224) plt.title("x") plt.scatter(x[:, 0], x[:, 1], alpha=.2) plt.xticks(np.arange(int(x_min[0]), int(x_max[0]), ticks[0]), np.arange(int(x_min[0]), int(x_max[0]), ticks[0])) plt.yticks(np.arange(int(x_min[1]), int(x_max[1]), ticks[1]), np.arange(int(x_min[1]), int(x_max[1]), ticks[1])) plt.suptitle( str(("epoch: ", epoch, "Train loss: ", ll_tot.item(), "Test loss: ", ll_test.item()))) plt.savefig("%s/%d.png" % (folder + toy, epoch)) plt.close(fig)
log_message = '[TEST] Iter {:04d} | Test Loss {:.6f} | NFE {:.0f}'.format(itr, test_loss, test_nfe) logger.info(log_message) if test_loss.item() < best_loss: best_loss = test_loss.item() utils.makedirs(args.save) torch.save({ 'args': args, 'state_dict': model.state_dict(), }, os.path.join(args.save, 'checkpt.pth')) model.train() if itr % args.viz_freq == 0: with torch.no_grad(): model.eval() p_samples = toy_data.inf_train_gen(args.data, batch_size=2000) sample_fn, density_fn = get_transforms(model) plt.figure(figsize=(9, 3)) visualize_transform( p_samples, torch.randn, standard_normal_logprob, transform=sample_fn, inverse_transform=density_fn, samples=True, npts=800, device=device ) fig_filename = os.path.join(args.save, 'figs', '{:04d}.pdf'.format(itr)) utils.makedirs(os.path.dirname(fig_filename)) plt.savefig(fig_filename) plt.close() model.train() end = time.time()
def loss_fn2(genFGen2, args, model): first_term_loss = compute_loss2(genFGen2, args, model) #print('') #print(first_term_loss) #import math #mu = torch.from_numpy(np.array([2.805741, -0.00889241], dtype="float32")).to(device) #S = torch.from_numpy(np.array([[pow(0.3442525,2), 0.0], [0.0, pow(0.35358343,2)]], dtype="float32")).to(device) #storeAll = torch.from_numpy(np.array(0.0, dtype="float32")).to(device) #toUse_storeAll = torch.distributions.MultivariateNormal(loc=mu, covariance_matrix=S) #for loopIndex_i in range(genFGen2.size()[0]): # storeAll += torch.exp(toUse_storeAll.log_prob(genFGen2[loopIndex_i:1 + loopIndex_i, :].squeeze(0))) #storeAll /= genFGen2.size()[0] #print(storeAll) #print('') #print('') #print(compute_loss2(mu.unsqueeze(0), args, model)) #print(torch.exp(toUse_storeAll.log_prob(mu))) #print('') #first_term_loss = storeAll xData = toy_data.inf_train_gen(args.data, batch_size=args.batch_size) xData = torch.from_numpy(xData).type(torch.float32).to(device) var2 = [] for i in genFGen2: var1 = [] for j in xData: new_stuff = torch.dist(i, j, 2) # this is a tensor var1.append(new_stuff.unsqueeze(0)) var1_tensor = torch.cat(var1) second_term_loss2 = torch.min(var1_tensor) / args.batch_size var2.append(second_term_loss2.unsqueeze(0)) var2_tensor = torch.cat(var2) second_term_loss = torch.mean(var2_tensor) / args.batch_size second_term_loss *= 100.0 print('') print(first_term_loss) print(second_term_loss) #third_term_loss = torch.from_numpy(np.array(0.0, dtype='float32')).to(device) #for i in range(args.batch_size): # for j in range(args.batch_size): # if i != j: # # third_term_loss += ((np.linalg.norm(genFGen3[i,:].cpu().detach().numpy()-genFGen3[j,:].cpu().detach().numpy())) / (np.linalg.norm(genFGen2[i,:].cpu().detach().numpy()-genFGen2[j,:].cpu().detach().numpy()))) # # # third_term_loss += ((torch.norm(genFGen3[i,:]-genFGen3[j,:], 2)) / (torch.norm(genFGen2[i,:]-genFGen2[j,:], 2))) # # third_term_loss += ((torch.norm(genFGen3[i,:]-genFGen3[j,:])) / (torch.norm(genFGen2[i,:]-genFGen2[j,:]))) # # # third_term_loss += ((torch.norm(genFGen3[i,:] - genFGen3[j,:])) / (torch.norm(genFGen2[i,:] - genFGen2[j,:]))) # third_term_loss += ((torch.dist(genFGen3[i, :], genFGen3[j, :], 2)) / (torch.dist(genFGen2[i, :], genFGen2[j, :], 2))) # third_term_loss /= (args.batch_size - 1) #third_term_loss /= args.batch_size ##third_term_loss *= 1000.0 #print(third_term_loss) #print('') print('') #asdfsfa #return first_term_loss + second_term_loss + third_term_loss return first_term_loss + second_term_loss
model = construct_discrete_model().to(device) model.load_state_dict(torch.load(args.checkpt)['state_dict']) else: model = build_model_tabular(args, 2).to(device) sd = torch.load(args.checkpt)['state_dict'] fixed_sd = {} for k, v in sd.items(): fixed_sd[k.replace('odefunc.odefunc', 'odefunc')] = v model.load_state_dict(fixed_sd) print(model) print("Number of trainable parameters: {}".format(count_parameters(model))) model.eval() p_samples = toy_data.inf_train_gen(args.data, batch_size=800**2) with torch.no_grad(): sample_fn, density_fn = get_transforms(model) plt.figure(figsize=(10, 10)) ax = ax = plt.gca() viz_flow.plt_samples(p_samples, ax, npts=800) plt.subplots_adjust(left=0, right=1, top=1, bottom=0) fig_filename = os.path.join(args.save, 'figs', 'true_samples.jpg') utils.makedirs(os.path.dirname(fig_filename)) plt.savefig(fig_filename) plt.close() plt.figure(figsize=(10, 10)) ax = ax = plt.gca()
def train_toy(toy, load=True, nb_step_dual=300, nb_steps=15, folder="", l1=1., nb_epoch=20000, pre_heating_epochs=10, nb_flow=3, cond_type="Coupling", emb_net=[150, 150, 150]): logger = utils.get_logger(logpath=os.path.join(folder, toy, 'logs'), filepath=os.path.abspath(__file__)) logger.info("Creating model...") device = "cpu" if not (torch.cuda.is_available()) else "cuda:0" nb_samp = 100 batch_size = 100 x_test = torch.tensor(toy_data.inf_train_gen(toy, batch_size=1000)).to(device) x = torch.tensor(toy_data.inf_train_gen(toy, batch_size=1000)).to(device) dim = x.shape[1] norm_type = "Affine" save_name = norm_type + str(emb_net) + str(nb_flow) solver = "CCParallel" int_net = [150, 150, 150] conditioner_type = cond_types[cond_type] conditioner_args = { "in_size": dim, "hidden": emb_net[:-1], "out_size": emb_net[-1] } if conditioner_type is DAGConditioner: conditioner_args['l1'] = l1 conditioner_args['gumble_T'] = .5 conditioner_args['nb_epoch_update'] = nb_step_dual conditioner_args["hot_encoding"] = True normalizer_type = norm_types[norm_type] if normalizer_type is MonotonicNormalizer: normalizer_args = { "integrand_net": int_net, "cond_size": emb_net[-1], "nb_steps": nb_steps, "solver": solver } else: normalizer_args = {} model = buildFCNormalizingFlow(nb_flow, conditioner_type, conditioner_args, normalizer_type, normalizer_args) opt = torch.optim.Adam(model.parameters(), 1e-4, weight_decay=1e-5) if load: logger.info("Loading model...") model.load_state_dict( torch.load(folder + toy + '/' + save_name + 'model.pt')) model.train() opt.load_state_dict( torch.load(folder + toy + '/' + save_name + 'ADAM.pt')) logger.info("Model loaded.") if True: for step in model.steps: step.conditioner.stoch_gate = True step.conditioner.noise_gate = False step.conditioner.gumble_T = .5 torch.autograd.set_detect_anomaly(True) for epoch in range(nb_epoch): loss_tot = 0 start = timer() for j in range(0, nb_samp, batch_size): cur_x = torch.tensor( toy_data.inf_train_gen(toy, batch_size=batch_size)).to(device) z, jac = model(cur_x) loss = model.loss(z, jac) loss_tot += loss.detach() if math.isnan(loss.item()): ll, z = model.compute_ll(cur_x) print(ll) print(z) print(ll.max(), z.max()) exit() opt.zero_grad() loss.backward(retain_graph=True) opt.step() model.step(epoch, loss_tot) end = timer() z, jac = model(x_test) ll = (model.z_log_density(z) + jac) ll_test = -ll.mean() dagness = max(model.DAGness()) logger.info( "epoch: {:d} - Train loss: {:4f} - Test loss: {:4f} - <<DAGness>>: {:4f} - Elapsed time per epoch {:4f} (seconds)" .format(epoch, loss_tot.item(), ll_test.item(), dagness, end - start)) if epoch % 100 == 0 and False: with torch.no_grad(): stoch_gate = model.getDag().stoch_gate noise_gate = model.getDag().noise_gate s_thresh = model.getDag().s_thresh model.getDag().stoch_gate = False model.getDag().noise_gate = False model.getDag().s_thresh = True for threshold in [.95, .1, .01, .0001, 1e-8]: model.set_h_threshold(threshold) # Valid loop z, jac = model(x_test) ll = (model.z_log_density(z) + jac) ll_test = -ll.mean().item() dagness = max(model.DAGness()).item() logger.info( "epoch: {:d} - Threshold: {:4f} - Valid log-likelihood: {:4f} - <<DAGness>>: {:4f}" .format(epoch, threshold, ll_test, dagness)) model.getDag().stoch_gate = stoch_gate model.getDag().noise_gate = noise_gate model.getDag().s_thresh = s_thresh model.set_h_threshold(0.) if epoch % 500 == 0: font = {'family': 'normal', 'weight': 'normal', 'size': 25} matplotlib.rc('font', **font) if toy in [ "2spirals-8gaussians", "4-2spirals-8gaussians", "8-2spirals-8gaussians", "2gaussians", "4gaussians", "2igaussians", "8gaussians" ] or True: def compute_ll(x): z, jac = model(x) ll = (model.z_log_density(z) + jac) return ll, z with torch.no_grad(): npts = 100 plt.figure(figsize=(12, 12)) gs = gridspec.GridSpec(2, 2, width_ratios=[3, 1], height_ratios=[3, 1]) ax = plt.subplot(gs[0]) qz_1, qz_2 = vf.plt_flow(compute_ll, ax, npts=npts, device=device) plt.subplot(gs[1]) plt.plot(qz_1, np.linspace(-4, 4, npts)) plt.ylabel('$x_2$', fontsize=25, rotation=-90, labelpad=20) plt.xticks([]) plt.subplot(gs[2]) plt.plot(np.linspace(-4, 4, npts), qz_2) plt.xlabel('$x_1$', fontsize=25) plt.yticks([]) plt.savefig("%s%s/flow_%s_%d.pdf" % (folder, toy, save_name, epoch)) torch.save(model.state_dict(), folder + toy + '/' + save_name + 'model.pt') torch.save(opt.state_dict(), folder + toy + '/' + save_name + 'ADAM.pt')
def loss_fn2(genFGen2, args, model): #print(list(model.parameters())) #logger.info(model) #logger.info("Number of trainable parameters: {}".format(count_parameters(model))) #with torch.no_grad(): # model.eval() # # p_samples = toy_data.inf_train_gen(args.data, batch_size=20000) # sample_fn, density_fn = model.inverse, model.forward # # plt.figure(figsize=(9, 3)) # visualize_transform(p_samples, torch.randn, standard_normal_logprob, transform=sample_fn, # inverse_transform=density_fn, samples=True, npts=400, device=device) # plt.show() # plt.close() # # plt.figure(figsize=(4, 2)) # visualize_transform2(p_samples, torch.randn, standard_normal_logprob, transform=sample_fn, # inverse_transform=density_fn, samples=True, npts=400, device=device) # plt.show() # plt.close() #xData = toy_data.inf_train_gen(args.data, batch_size=args.batch_size) #xData = torch.from_numpy(xData).type(torch.float32).to(device) #second_term_loss = torch.from_numpy(np.array(0.0, dtype='float32')).to(device) #for i in genFGen2: # for j in xData: # # second_term_loss += np.linalg.norm(i.cpu().detach().numpy()-j.cpu().detach().numpy()) # # # second_term_loss += torch.norm(i-j, 2) # # second_term_loss += torch.norm(i-j) # # # second_term_loss += torch.dist(i.type(torch.float64), j.type(torch.float64), 2) # second_term_loss += torch.dist(i, j, 2) # second_term_loss /= args.batch_size # #second_term_loss *= 0.1 #second_term_loss /= args.batch_size #print('') #print(second_term_loss) #second_term_loss = torch.from_numpy(np.array([], dtype='float32')).to(device) #for i in genFGen2: # second_term_lossTwo = torch.from_numpy(np.array([], dtype='float32')).to(device) # for j in xData: # second_term_lossTwo = torch.cat((second_term_lossTwo, torch.dist(i, j, 2)), 0) # second_term_loss = torch.cat((second_term_loss, torch.mean(second_term_lossTwo, 0)), 0) # #second_term_loss *= 0.1 #second_term_loss = torch.mean(second_term_loss) #print('') #print(second_term_loss) #genFGen3 = torch.randn((args.batch_size, 2)).to(device) #first_term_loss = compute_loss2(genFGen2, args, model, beta=beta) #print('') #print(first_term_loss) import math mu = torch.from_numpy(np.array([2.805741, -0.00889241], dtype="float32")).to(device) S = torch.from_numpy(np.array([[pow(0.3442525,2), 0.0], [0.0, pow(0.35358343,2)]], dtype="float32")).to(device) storeAll = torch.from_numpy(np.array(0.0, dtype="float32")).to(device) for loopIndex_i in range(genFGen2.size()[0]): toUse_storeAll = torch.distributions.MultivariateNormal(loc=mu, covariance_matrix=S) storeAll += toUse_storeAll.log_prob(genFGen2[loopIndex_i:1+loopIndex_i, :].squeeze(0)) storeAll /= genFGen2.size()[0] #print(storeAll) #print('') first_term_loss = storeAll xData = toy_data.inf_train_gen(args.data, batch_size=args.batch_size) xData = torch.from_numpy(xData).type(torch.float32).to(device) var2 = [] for i in genFGen2: var1 = [] for j in xData: new_stuff = torch.dist(i, j, 2) # this is a tensor var1.append(new_stuff.unsqueeze(0)) var1_tensor = torch.cat(var1) second_term_loss2 = torch.min(var1_tensor) / args.batch_size var2.append(second_term_loss2.unsqueeze(0)) var2_tensor = torch.cat(var2) second_term_loss = torch.mean(var2_tensor) / args.batch_size first_term_loss = torch.exp(first_term_loss) print('') print(first_term_loss) print(second_term_loss) #third_term_loss = torch.from_numpy(np.array(0.0, dtype='float32')).to(device) #for i in range(args.batch_size): # for j in range(args.batch_size): # if i != j: # # third_term_loss += ((np.linalg.norm(genFGen3[i,:].cpu().detach().numpy()-genFGen3[j,:].cpu().detach().numpy())) / (np.linalg.norm(genFGen2[i,:].cpu().detach().numpy()-genFGen2[j,:].cpu().detach().numpy()))) # # # third_term_loss += ((torch.norm(genFGen3[i,:]-genFGen3[j,:], 2)) / (torch.norm(genFGen2[i,:]-genFGen2[j,:], 2))) # # third_term_loss += ((torch.norm(genFGen3[i,:]-genFGen3[j,:])) / (torch.norm(genFGen2[i,:]-genFGen2[j,:]))) # # # third_term_loss += ((torch.norm(genFGen3[i,:] - genFGen3[j,:])) / (torch.norm(genFGen2[i,:] - genFGen2[j,:]))) # third_term_loss += ((torch.dist(genFGen3[i, :], genFGen3[j, :], 2)) / (torch.dist(genFGen2[i, :], genFGen2[j, :], 2))) # third_term_loss /= (args.batch_size - 1) #third_term_loss /= args.batch_size ##third_term_loss *= 1000.0 #print(third_term_loss) #print('') print('') #asdfsfa #return first_term_loss + second_term_loss + third_term_loss return first_term_loss + second_term_loss
logger.info("-------------------------") logger.info(str(optim)) # optimizer info logger.info("data={:} batch_size={:} gpu={:}".format( args.data, args.batch_size, args.gpu)) logger.info("maxIters={:} val_freq={:} viz_freq={:}".format( args.niters, args.val_freq, args.viz_freq)) logger.info("saveLocation = {:}".format(args.save)) logger.info("-------------------------\n") end = time.time() best_loss = float('inf') bestParams = None # setup data [nSamples, d] # use one batch as the entire data set x0 = toy_data.inf_train_gen(args.data, batch_size=args.batch_size) x0 = cvt(torch.from_numpy(x0)) x0val = toy_data.inf_train_gen(args.data, batch_size=args.val_batch_size) x0val = cvt(torch.from_numpy(x0val)) log_msg = ( '{:5s} {:6s} {:9s} {:9s} {:9s} {:9s} {:9s} {:9s} {:9s} {:9s} ' .format('iter', ' time', 'loss', 'L (L_2)', 'C (loss)', 'R (HJB)', 'valLoss', 'valL', 'valC', 'valR')) logger.info(log_msg) time_meter = utils.AverageMeter() net.train() for itr in range(1, args.niters + 1):