From 71164326cf0bf447cce4fd23be2cf661425cb0fb Mon Sep 17 00:00:00 2001 From: Johannes Buchner Date: Thu, 2 May 2019 20:14:37 +0200 Subject: [PATCH 1/3] iterative sampling + training (multitrainer) --- examples/mcmc/himmelblau.py | 7 +- examples/mcmc/rosenbrock.py | 7 +- nnest/mcmc.py | 15 +- nnest/multitrainer.py | 536 ++++++++++++++++++++++++++++++++++++ nnest/sampler.py | 2 +- nnest/trainer.py | 75 ++--- setup.py | 2 +- 7 files changed, 601 insertions(+), 43 deletions(-) create mode 100644 nnest/multitrainer.py diff --git a/examples/mcmc/himmelblau.py b/examples/mcmc/himmelblau.py index 1a5bf52..baaa854 100644 --- a/examples/mcmc/himmelblau.py +++ b/examples/mcmc/himmelblau.py @@ -17,11 +17,12 @@ def loglike(z): return - (z1**2 + z2 - 11.)**2 - (z1 + z2**2 - 7.)**2 def transform(x): - return 5. * x + return 500. * x sampler = MCMCSampler(args.x_dim, loglike, transform=transform, log_dir=args.log_dir, hidden_dim=args.hidden_dim, num_layers=args.num_layers, num_blocks=args.num_blocks, num_slow=args.num_slow) - sampler.run(train_iters=args.train_iters, mcmc_steps=args.mcmc_steps, single_thin=10) + sampler.run(train_iters=args.train_iters, mcmc_steps=args.mcmc_steps, single_thin=10, + bootstrap_iters=args.burnin_iters) if __name__ == '__main__': @@ -32,6 +33,8 @@ def transform(x): help="Dimensionality") parser.add_argument('--train_iters', type=int, default=200, help="number of train iters") + parser.add_argument('--burnin_iters', type=int, default=1, + help="number of iters for finding good training spot") parser.add_argument("--mcmc_steps", type=int, default=10000) parser.add_argument('--hidden_dim', type=int, default=128) parser.add_argument('--num_layers', type=int, default=2) diff --git a/examples/mcmc/rosenbrock.py b/examples/mcmc/rosenbrock.py index ab75bee..2758a4e 100644 --- a/examples/mcmc/rosenbrock.py +++ b/examples/mcmc/rosenbrock.py @@ -16,11 +16,12 @@ def loglike(z): return np.array([-sum(100.0 * (x[1:] - x[:-1] ** 2.0) ** 2.0 + (1 - x[:-1]) ** 2.0) for x in z]) def transform(x): - return 5. * x + return 500. * x - 100 sampler = MCMCSampler(args.x_dim, loglike, transform=transform, log_dir=args.log_dir, hidden_dim=args.hidden_dim, num_layers=args.num_layers, num_blocks=args.num_blocks, num_slow=args.num_slow) - sampler.run(train_iters=args.train_iters, mcmc_steps=args.mcmc_steps, single_thin=10) + sampler.run(train_iters=args.train_iters, mcmc_steps=args.mcmc_steps, single_thin=10, + bootstrap_iters=args.burnin_iters, bootstrap_mcmc_steps=5000) if __name__ == '__main__': @@ -30,6 +31,8 @@ def transform(x): help="Dimensionality") parser.add_argument('--train_iters', type=int, default=100, help="number of train iters") + parser.add_argument('--burnin_iters', type=int, default=1, + help="number of iters for finding good training spot") parser.add_argument("--mcmc_steps", type=int, default=10000) parser.add_argument('--hidden_dim', type=int, default=128) parser.add_argument('--num_layers', type=int, default=1) diff --git a/nnest/mcmc.py b/nnest/mcmc.py index bbf6b52..9e6fa3b 100644 --- a/nnest/mcmc.py +++ b/nnest/mcmc.py @@ -109,14 +109,15 @@ def run( bootstrap_batch_size=5, alpha=0, single_thin=1, - ignore_rows=0.3): + ignore_rows=0.0): if alpha == 0.0: alpha = 1 / self.x_dim ** 0.5 if self.log: self.logger.info('Alpha [%5.4f]' % (alpha)) - + + next_start = None for t in range(bootstrap_iters): if t == 0: @@ -128,9 +129,13 @@ def run( else: samples, likes, latent, scale, nc = self.trainer.sample( loglike=self.loglike, transform=transform, - mcmc_steps=bootstrap_mcmc_steps, alpha=alpha, dynamic=False, show_progress=True) + mcmc_steps=bootstrap_mcmc_steps, + alpha=alpha, dynamic=False, show_progress=True, + init_x=next_start) + next_start = samples[:,-1,:] samples = transform(samples) self._chain_stats(samples) + self.logger.info('Last sample: %s' % (next_start)) loglikes = -np.array(likes) weights = np.ones(loglikes.shape) mc = MCSamples(samples=[samples[0]], weights=[weights[0]], loglikes=[loglikes[0]], @@ -141,11 +146,15 @@ def run( mean = np.mean(samples, axis=0) std = np.std(samples, axis=0) samples = (samples - mean) / std + self.logger.info('Training for bootstrap iter %d with %d samples' % (t, len(samples))) + # Forget the current network + self.trainer.init_network() self.trainer.train(samples, max_iters=train_iters, noise=-1) def transform(x): return x * std + mean + self.logger.info('Sampling...') samples, likes, latent, scale, nc = self.trainer.sample( loglike=self.loglike, transform=transform, mcmc_steps=mcmc_steps, alpha=alpha, dynamic=False, show_progress=True, diff --git a/nnest/multitrainer.py b/nnest/multitrainer.py new file mode 100644 index 0000000..4481238 --- /dev/null +++ b/nnest/multitrainer.py @@ -0,0 +1,536 @@ +""" +.. module:: trainer + :synopsis: Train the flow and perform sampling in latent space +.. moduleauthor:: Adam Moss +""" + +from __future__ import print_function +from __future__ import division + +import os +import time + +import torch +from torch.utils.data import DataLoader +import numpy as np +from tensorboardX import SummaryWriter +from sklearn.model_selection import train_test_split +import scipy.spatial +from tqdm import tqdm + +import matplotlib + +matplotlib.use('Agg') +import matplotlib.pyplot as plt + +from nnest.networks import SingleSpeed, FastSlow, BatchNormFlow +from nnest.utils.logger import create_logger + + +class MultiTrainer(object): + + def __init__(self, + xdim, + ndim, + nslow=0, + batch_size=100, + flow='nvp', + num_blocks=5, + num_layers=2, + oversample_rate=-1, + train=True, + load_model='', + log_dir='logs', + use_gpu=False, + log=True + ): + + self.device = torch.device( + 'cuda' if use_gpu and torch.cuda.is_available() else 'cpu') + + self.x_dim = xdim + self.z_dim = xdim + + self.batch_size = batch_size + self.total_iters = 0 + + assert xdim > nslow + nfast = xdim - nslow + self.nslow = nslow + + if oversample_rate > 0: + self.oversample_rate = oversample_rate + else: + self.oversample_rate = nfast / xdim + + def init_network(): + if flow.lower() == 'nvp': + if nslow > 0: + self.netG = FastSlow(nfast, nslow, ndim, num_blocks, num_layers) + else: + self.netG = SingleSpeed(xdim, ndim, num_blocks, num_layers) + else: + raise NotImplementedError + self.nparams = sum(p.numel() for p in self.netG.parameters()) + + if train and not load_model: + if log_dir is not None: + self.path = log_dir + if not os.path.exists(os.path.join(self.path, 'models')): + os.makedirs(os.path.join(self.path, 'models')) + if not os.path.exists(os.path.join(self.path, 'data')): + os.makedirs(os.path.join(self.path, 'data')) + if not os.path.exists(os.path.join(self.path, 'chains')): + os.makedirs(os.path.join(self.path, 'chains')) + if not os.path.exists(os.path.join(self.path, 'plots')): + os.makedirs(os.path.join(self.path, 'plots')) + else: + self.path = None + else: + self.path = os.path.join(log_dir, load_model) + self.netG.load_state_dict(torch.load( + os.path.join(self.path, 'models', 'netG.pt') + )) + self.optimizer = torch.optim.Adam( + self.netG.parameters(), lr=0.0001, weight_decay=1e-6) + + if self.path is not None: + self.writer = SummaryWriter(self.path) + + self.init_network = init_network + self.init_network() + print(self.netG) + + self.logger = create_logger(__name__) + self.log = log + + def train( + self, + samples, + max_iters=5000, + log_interval=50, + save_interval=50, + noise=0.0, + validation_fraction=0.05): + + start_time = time.time() + + if self.path: + fig, ax = plt.subplots() + ax.scatter(samples[:, 0], samples[:, 1]) + self.writer.add_figure('originals', fig, self.total_iters) + np.save( + os.path.join(self.path, 'data', 'originals.npy'), + samples) + + if noise < 0: + # compute distance to nearest neighbor + kdt = scipy.spatial.cKDTree(samples) + dists, neighs = kdt.query(samples, 2) + training_noise = .2 * np.mean(dists) + else: + training_noise = noise + + if self.log: + self.logger.info('Number of training samples [%d] for [%d] variables' % (samples.shape[0], self.nparams)) + self.logger.info('Training noise [%5.4f]' % training_noise) + + X_train, X_valid = train_test_split( + samples, test_size=validation_fraction) + + train_tensor = torch.from_numpy(X_train.astype(np.float32)) + train_dataset = torch.utils.data.TensorDataset(train_tensor) + train_loader = torch.utils.data.DataLoader( + train_dataset, batch_size=self.batch_size, shuffle=True) + + valid_tensor = torch.from_numpy(X_valid.astype(np.float32)) + valid_dataset = torch.utils.data.TensorDataset(valid_tensor) + valid_loader = torch.utils.data.DataLoader( + valid_dataset, batch_size=X_valid.shape[0], shuffle=False, drop_last=True) + + best_validation_loss = float('inf') + best_validation_epoch = 0 + best_model = self.netG.state_dict() + self.netGs = [] + self.netGs_accepts = [] + self.netGs_samples = [] + self.netGs_init_x = [] + + for epoch in range(1, max_iters + 1): + + self.total_iters += 1 + + train_loss = self._train( + epoch, self.netG, train_loader, noise=training_noise) + validation_loss = self._validate(epoch, self.netG, valid_loader) + + stored = ' ' + # remember only if: + # 1) much better + # 2) when performance is improving drastically at each epoch, avoid every one of them. + if validation_loss < best_validation_loss - 1.0 and epoch >= best_validation_epoch + save_interval: + best_validation_epoch = epoch + best_validation_loss = validation_loss + best_model = self.netG.state_dict() + self.netGs.append(self.netG.state_dict()) + self.netGs_accepts.append(0) + self.netGs_samples.append(0) + self.netGs_init_x.append(None) + stored = '**' + + if epoch == 1 or epoch % log_interval == 0: + print( + 'Epoch: {} validation loss: {:6.4f} {}'.format( + epoch, validation_loss, stored)) + + #if epoch % save_interval == 0: + + if self.path: + self.writer.add_scalar('loss', validation_loss, self.total_iters) + if epoch % save_interval == 0: + torch.save( + self.netG.state_dict(), + os.path.join(self.path, 'models', 'netG.pt%d' % epoch) + ) + self._train_plot(self.netG, samples) + self.netG.load_state_dict(best_model) + + def choose_netG(self): + # compute accept probabilities + print("Networks:", ' '.join(['%3d/%3d' % (A, S) for A, S in zip(self.netGs_accepts, self.netGs_samples)])) + # give initial acceptance probabilities of 1% + probs = np.array([(A+1)/(S+1) for A, S in zip(self.netGs_accepts, self.netGs_samples)]) + probs /= probs.sum() + i = np.random.choice(np.arange(len(probs)), p=probs) + self.netG.load_state_dict(self.netGs[i]) + #print("Network %d: %d/%d" % (i, self.netGs_accepts[i], self.netGs_samples[i])) + return i + + def sample( + self, + mcmc_steps=20, + alpha=1.0, + dynamic=True, + batch_size=1, + loglike=None, + init_x=None, + logl=None, + loglstar=None, + transform=None, + show_progress=False, + plot=False, + out_chain=None, + max_prior=None, + max_start_tries=100): + + allsamples = [] + alllatent = [] + alllikes = [] + allncall = 0 + nsamples = 20 + scales = [alpha * 1.0 for _ in self.netGs_init_x] + while len(allsamples) < mcmc_steps: + neti = self.choose_netG() + init_x = self.netGs_init_x[neti] + scale = scales[neti] + + samples, likes, latent, scale, ncall, accept = self.subsample( + mcmc_steps=nsamples, alpha=scale, dynamic=dynamic, + loglike=loglike, init_x=init_x, logl=logl, loglstar=loglstar, + transform=transform, show_progress=False, plot=plot, + out_chain=out_chain, max_prior=max_prior, max_start_tries=max_start_tries, + ) + + self.netGs_accepts[neti] += accept.min() + self.netGs_samples[neti] += nsamples + self.netGs_init_x[neti] = samples[-1] + scales[neti] = scale + + allsamples += samples + alllikes += likes + alllatent += latent + allncall = allncall + ncall + self.logger.info('scales: %s' % scales) + + # Transpose so shape is (chain_num, iteration, dim) + samples = np.transpose(np.array(allsamples), axes=[1, 0, 2]) + likes = np.transpose(np.array(alllikes), axes=[1, 0]) + latent = np.transpose(np.array(alllatent), axes=[1, 0, 2]) + ncall = allncall + + if self.path and plot: + cmap = plt.cm.jet + cmap.set_under('w', 1) + fig, ax = plt.subplots() + ax.hist2d(samples[:, -1, 0], samples[:, -1, 1], + bins=200, cmap=cmap, vmin=1, alpha=0.2) + if self.writer is not None: + self.writer.add_figure('chain', fig, self.total_iters) + + if out_chain is not None: + for ib in range(batch_size): + files[ib].close() + + return samples, likes, latent, scale, ncall + + + def subsample( + self, + mcmc_steps=20, + alpha=1.0, + dynamic=True, + batch_size=1, + loglike=None, + init_x=None, + logl=None, + loglstar=None, + transform=None, + show_progress=False, + plot=False, + out_chain=None, + max_prior=None, + max_start_tries=100): + + self.netG.eval() + + samples = [] + latent = [] + likes = [] + + if transform is None: + def transform(x): return x + + if init_x is not None: + batch_size = init_x.shape[0] + z, _ = self.netG(torch.from_numpy(init_x).float().to(self.device)) + z = z.detach() + # Add the backward version of x rather than init_x due to numerical precision + x, _ = self.netG(z, mode='inverse') + x = x.detach().cpu().numpy() + if logl is None: + logl = loglike(transform(x)) + else: + if logl is None: + for i in range(max_start_tries): + z = torch.randn(batch_size, self.z_dim, device=self.device) + x, _ = self.netG(z, mode='inverse') + x = x.detach().cpu().numpy() + logl = loglike(transform(x)) + if np.all(logl > -1e30): + break + if i == max_start_tries - 1: + raise Exception('Could not find starting value') + else: + z = torch.randn(batch_size, self.z_dim, device=self.device) + x, _ = self.netG(z, mode='inverse') + x = x.detach().cpu().numpy() + logl = loglike(transform(x)) + + samples.append(x) + likes.append(logl) + + iters = range(mcmc_steps) + if show_progress: + iters = tqdm(iters) + + scale = alpha + accept = 0 + reject = 0 + ncall = 0 + + if out_chain is not None: + if batch_size == 1: + files = [open(out_chain + '.txt', 'w')] + else: + files = [open(out_chain + '_%s.txt' % (ib + 1), 'w') for ib in range(batch_size)] + + for i in iters: + + dz = torch.randn_like(z) * scale + if self.nslow > 0 and np.random.uniform() < self.oversample_rate: + fast = True + dz[:, 0:self.nslow] = 0.0 + else: + fast = False + z_prime = z + dz + + # Jacobian is det d f^{-1} (z)/dz + x, log_det_J = self.netG(z, mode='inverse') + x_prime, log_det_J_prime = self.netG(z_prime, mode='inverse') + x = x.detach().cpu().numpy() + x_prime = x_prime.detach().cpu().numpy() + delta_log_det_J = (log_det_J_prime - log_det_J).detach() + log_ratio_1 = delta_log_det_J.squeeze(dim=1) + + # Check not out of prior range + if max_prior is not None: + prior = np.logical_or( + np.abs(x) > max_prior, + np.abs(x_prime) > max_prior) + idx = np.where([np.any(p) for p in prior]) + log_ratio_1[idx] = -np.inf + + rnd_u = torch.rand(log_ratio_1.shape, device=self.device) + ratio = log_ratio_1.exp().clamp(max=1) + mask = (rnd_u < ratio).int() + logl_prime = np.full(batch_size, logl) + + # Only evaluate likelihood if prior and volume is accepted + if loglike is not None and transform is not None: + for idx, im in enumerate(mask): + if im: + if not fast: + ncall += 1 + lp = loglike(transform(x_prime[idx])) + if loglstar is not None: + if np.isfinite(lp) and lp >= loglstar: + logl_prime[idx] = lp + else: + mask[idx] = 0 + else: + if lp >= logl[idx]: + logl_prime[idx] = lp + elif rnd_u[idx].cpu().numpy() < np.clip(np.exp(lp - logl[idx]), 0, 1): + logl_prime[idx] = lp + else: + mask[idx] = 0 + + accept += torch.sum(mask).cpu().numpy() + reject += batch_size - torch.sum(mask).cpu().numpy() + + if dynamic: + if accept > reject: + scale *= np.exp(1. / accept) + if accept < reject: + scale /= np.exp(1. / reject) + + m = mask[:, None].float() + z = (z_prime * m + z * (1 - m)).detach() + + m = mask + logl = logl_prime * m.cpu().numpy() + logl * (1 - m.cpu().numpy()) + + x, _ = self.netG(z, mode='inverse') + x = x.detach().cpu().numpy() + samples.append(x) + likes.append(logl) + latent.append(z.cpu().numpy()) + + if out_chain is not None: + v = transform(x) + for ib in range(batch_size): + files[ib].write("%.5E " % 1) + files[ib].write("%.5E " % -logl[ib]) + files[ib].write(" ".join(["%.5E" % vi for vi in v[ib]])) + files[ib].write("\n") + files[ib].flush() + + return samples, likes, latent, scale, ncall, accept + + def _jacobian(self, z): + """ Calculate det d f^{-1} (z)/dz + """ + z.requires_grad_(True) + x, _ = self.netG(z, mode='inverse') + J = torch.stack([torch.autograd.grad(outputs=x[:, i], inputs=z, retain_graph=True, + grad_outputs=torch.ones(z.shape[0]))[0] for i in + range(x.shape[1])]).permute(1, 0, 2) + return torch.stack([torch.log(torch.abs(torch.det(J[i, :, :]))) + for i in range(x.shape[0])]) + + def _train(self, epoch, model, loader, noise=0.0): + + model.train() + train_loss = 0 + + for batch_idx, data in enumerate(loader): + if isinstance(data, list): + if len(data) > 1: + cond_data = data[1].float() + cond_data = cond_data.to(self.device) + else: + cond_data = None + data = data[0] + data = (data + noise * torch.randn_like(data)).to(self.device) + self.optimizer.zero_grad() + loss = -model.log_probs(data, cond_data).mean() + train_loss += loss.item() + loss.backward() + self.optimizer.step() + + for module in model.modules(): + if isinstance(module, BatchNormFlow): + module.momentum = 0 + + with torch.no_grad(): + model(loader.dataset.tensors[0].to(data.device)) + + for module in model.modules(): + if isinstance(module, BatchNormFlow): + module.momentum = 1 + + return train_loss / len(loader.dataset) + + def _validate(self, epoch, model, loader): + + model.eval() + val_loss = 0 + cond_data = None + + for batch_idx, data in enumerate(loader): + if isinstance(data, list): + if len(data) > 1: + cond_data = data[1].float() + cond_data = cond_data.to(self.device) + data = data[0] + data = data.to(self.device) + with torch.no_grad(): + # sum up batch loss + val_loss += -model.log_probs(data, cond_data).sum().item() + + return val_loss / len(loader.dataset) + + def _train_plot(self, model, dataset, plot_grid=True): + + model.eval() + with torch.no_grad(): + x_synth = model.sample(dataset.size).detach().cpu().numpy() + x_data = torch.from_numpy(dataset).float() + z, _ = model(x_data) + z = z.detach().cpu().numpy() + if plot_grid: + grid = [] + for x in np.linspace(np.min(dataset[:, 0]), np.max(dataset[:, 0]), 10): + for y in np.linspace(np.min(dataset[:, 1]), np.max(dataset[:, 1]), 5000): + grid.append([x, y]) + for y in np.linspace(np.min(dataset[:, 1]), np.max(dataset[:, 1]), 10): + for x in np.linspace(np.min(dataset[:, 0]), np.max(dataset[:, 0]), 5000): + grid.append([x, y]) + grid = np.array(grid) + x_grid = torch.from_numpy(grid).float() + z_grid, _ = model(x_grid) + z_grid = z_grid.detach().cpu().numpy() + + if self.writer is not None: + fig, ax = plt.subplots(2, figsize=(5, 10)) + ax[0].scatter(dataset[:, 0], dataset[:, 1], c=dataset[:, 0], s=4) + ax[1].scatter(z[:, 0], z[:, 1], c=dataset[:, 0], s=4) + self.writer.add_figure('latent', fig, self.total_iters) + + if self.path: + fig, ax = plt.subplots(1, 3, figsize=(10, 4)) + if plot_grid: + ax[0].scatter(grid[:, 0], grid[:, 1], c=grid[:, 0], marker='.', s=1, linewidths=0) + ax[0].scatter(dataset[:, 0], dataset[:, 1], s=4) + ax[0].set_title('Real data') + if plot_grid: + ax[1].scatter(z_grid[:, 0], z_grid[:, 1], c=grid[:, 0], marker='.', s=1, linewidths=0) + ax[1].scatter(z[:, 0], z[:, 1], s=4) + ax[1].set_title('Latent data') + ax[1].set_xlim([-3, 3]) + ax[1].set_ylim([-3, 3]) + ax[2].scatter(x_synth[:, 0], x_synth[:, 1], s=2) + ax[2].set_title('Synthetic data') + plt.tight_layout() + plt.savefig(os.path.join(self.path, 'plots', 'plot_%s.png' % self.total_iters)) + plt.close() diff --git a/nnest/sampler.py b/nnest/sampler.py index 5c37ef7..5ec6240 100755 --- a/nnest/sampler.py +++ b/nnest/sampler.py @@ -12,7 +12,7 @@ import numpy as np -from nnest.trainer import Trainer +from nnest.multitrainer import MultiTrainer as Trainer from nnest.utils.logger import create_logger, make_run_dir from nnest.utils.evaluation import acceptance_rate, effective_sample_size, mean_jump_distance diff --git a/nnest/trainer.py b/nnest/trainer.py index d8c7451..4da6f45 100644 --- a/nnest/trainer.py +++ b/nnest/trainer.py @@ -62,40 +62,44 @@ def __init__(self, self.oversample_rate = oversample_rate else: self.oversample_rate = nfast / xdim - - if flow.lower() == 'nvp': - if nslow > 0: - self.netG = FastSlow(nfast, nslow, ndim, num_blocks, num_layers) + + def init_network(): + if flow.lower() == 'nvp': + if nslow > 0: + self.netG = FastSlow(nfast, nslow, ndim, num_blocks, num_layers) + else: + self.netG = SingleSpeed(xdim, ndim, num_blocks, num_layers) else: - self.netG = SingleSpeed(xdim, ndim, num_blocks, num_layers) - else: - raise NotImplementedError - - if train and not load_model: - if log_dir is not None: - self.path = log_dir - if not os.path.exists(os.path.join(self.path, 'models')): - os.makedirs(os.path.join(self.path, 'models')) - if not os.path.exists(os.path.join(self.path, 'data')): - os.makedirs(os.path.join(self.path, 'data')) - if not os.path.exists(os.path.join(self.path, 'chains')): - os.makedirs(os.path.join(self.path, 'chains')) - if not os.path.exists(os.path.join(self.path, 'plots')): - os.makedirs(os.path.join(self.path, 'plots')) + raise NotImplementedError + self.nparams = sum(p.numel() for p in self.netG.parameters()) + + if train and not load_model: + if log_dir is not None: + self.path = log_dir + if not os.path.exists(os.path.join(self.path, 'models')): + os.makedirs(os.path.join(self.path, 'models')) + if not os.path.exists(os.path.join(self.path, 'data')): + os.makedirs(os.path.join(self.path, 'data')) + if not os.path.exists(os.path.join(self.path, 'chains')): + os.makedirs(os.path.join(self.path, 'chains')) + if not os.path.exists(os.path.join(self.path, 'plots')): + os.makedirs(os.path.join(self.path, 'plots')) + else: + self.path = None else: - self.path = None - else: - self.path = os.path.join(log_dir, load_model) - self.netG.load_state_dict(torch.load( - os.path.join(self.path, 'models', 'netG.pt') - )) - - self.optimizer = torch.optim.Adam( - self.netG.parameters(), lr=0.0001, weight_decay=1e-6) - - if self.path is not None: - print(self.netG) - self.writer = SummaryWriter(self.path) + self.path = os.path.join(log_dir, load_model) + self.netG.load_state_dict(torch.load( + os.path.join(self.path, 'models', 'netG.pt') + )) + self.optimizer = torch.optim.Adam( + self.netG.parameters(), lr=0.0001, weight_decay=1e-6) + + if self.path is not None: + self.writer = SummaryWriter(self.path) + + self.init_network = init_network + self.init_network() + print(self.netG) self.logger = create_logger(__name__) self.log = log @@ -127,7 +131,7 @@ def train( training_noise = noise if self.log: - self.logger.info('Number of training samples [%d]' % samples.shape[0]) + self.logger.info('Number of training samples [%d] for [%d] variables' % (samples.shape[0], self.nparams)) self.logger.info('Training noise [%5.4f]' % training_noise) X_train, X_valid = train_test_split( @@ -145,7 +149,7 @@ def train( best_validation_loss = float('inf') best_validation_epoch = 0 - best_model = self.netG + best_model = self.netG.state_dict() for epoch in range(1, max_iters + 1): @@ -158,6 +162,7 @@ def train( if validation_loss < best_validation_loss: best_validation_epoch = epoch best_validation_loss = validation_loss + best_model = self.netG.state_dict() if epoch == 1 or epoch % log_interval == 0: print( @@ -173,6 +178,8 @@ def train( ) self._train_plot(self.netG, samples) + self.netG.load_state_dict(best_model) + def sample( self, mcmc_steps=20, diff --git a/setup.py b/setup.py index 62e2aaa..a6a5ae8 100644 --- a/setup.py +++ b/setup.py @@ -3,5 +3,5 @@ setup( name='nnest', version='0.1.0', - packages=['nnest'], + packages=['nnest', 'nnest.utils'], ) From e16a5519f12282e46aac2a9d709bec8bf5d8d1f0 Mon Sep 17 00:00:00 2001 From: Johannes Buchner Date: Fri, 3 May 2019 09:11:06 +0200 Subject: [PATCH 2/3] keep multiple networks, sample proportionally to their efficiency --- examples/mcmc/rosenbrock.py | 10 ++- nnest/mcmc.py | 12 +++- nnest/multitrainer.py | 140 +++++------------------------------- nnest/networks.py | 4 +- nnest/trainer.py | 2 +- 5 files changed, 40 insertions(+), 128 deletions(-) diff --git a/examples/mcmc/rosenbrock.py b/examples/mcmc/rosenbrock.py index 2758a4e..b5c4a97 100644 --- a/examples/mcmc/rosenbrock.py +++ b/examples/mcmc/rosenbrock.py @@ -11,9 +11,15 @@ def main(args): from nnest.mcmc import MCMCSampler os.environ['CUDA_VISIBLE_DEVICES'] = '' - - def loglike(z): + + def loglike_orig(z): return np.array([-sum(100.0 * (x[1:] - x[:-1] ** 2.0) ** 2.0 + (1 - x[:-1]) ** 2.0) for x in z]) + + # pairwise rosenbrocks, multiplied together + assert args.x_dim % 2 == 0 + + def loglike(z): + return np.array([-sum(100.0 * (x[1::2] - x[::2] ** 2.0) ** 2.0 + (1 - x[::2]) ** 2.0) * 2. / len(x) for x in z]) def transform(x): return 500. * x - 100 diff --git a/nnest/mcmc.py b/nnest/mcmc.py index 9e6fa3b..e61a0fd 100644 --- a/nnest/mcmc.py +++ b/nnest/mcmc.py @@ -117,6 +117,7 @@ def run( if self.log: self.logger.info('Alpha [%5.4f]' % (alpha)) + allsamples = [] next_start = None for t in range(bootstrap_iters): @@ -133,9 +134,10 @@ def run( alpha=alpha, dynamic=False, show_progress=True, init_x=next_start) next_start = samples[:,-1,:] + #next_start = None samples = transform(samples) self._chain_stats(samples) - self.logger.info('Last sample: %s' % (next_start)) + #self.logger.info('Last sample: %s' % (next_start)) loglikes = -np.array(likes) weights = np.ones(loglikes.shape) mc = MCSamples(samples=[samples[0]], weights=[weights[0]], loglikes=[loglikes[0]], @@ -146,6 +148,12 @@ def run( mean = np.mean(samples, axis=0) std = np.std(samples, axis=0) samples = (samples - mean) / std + self.logger.info('%d new samples' % (len(samples))) + if t > 0: + allsamples = np.vstack((allsamples, samples)) + samples = allsamples[int(len(allsamples)//3):,:] + else: + allsamples = samples self.logger.info('Training for bootstrap iter %d with %d samples' % (t, len(samples))) # Forget the current network self.trainer.init_network() @@ -156,7 +164,7 @@ def transform(x): self.logger.info('Sampling...') samples, likes, latent, scale, nc = self.trainer.sample( - loglike=self.loglike, transform=transform, + loglike=self.loglike, transform=transform, init_x=next_start, mcmc_steps=mcmc_steps, alpha=alpha, dynamic=False, show_progress=True, out_chain=os.path.join(self.logs['chains'], 'chain')) samples = transform(samples) diff --git a/nnest/multitrainer.py b/nnest/multitrainer.py index 4481238..f0c4a37 100644 --- a/nnest/multitrainer.py +++ b/nnest/multitrainer.py @@ -25,9 +25,10 @@ from nnest.networks import SingleSpeed, FastSlow, BatchNormFlow from nnest.utils.logger import create_logger +from nnest.trainer import Trainer -class MultiTrainer(object): +class MultiTrainer(Trainer): def __init__(self, xdim, @@ -127,7 +128,7 @@ def train( # compute distance to nearest neighbor kdt = scipy.spatial.cKDTree(samples) dists, neighs = kdt.query(samples, 2) - training_noise = .2 * np.mean(dists) + training_noise = 0.5 * np.mean(dists) / self.x_dim else: training_noise = noise @@ -192,14 +193,16 @@ def train( self.netG.state_dict(), os.path.join(self.path, 'models', 'netG.pt%d' % epoch) ) - self._train_plot(self.netG, samples) + if self.x_dim == 2: + self._train_plot(self.netG, samples) self.netG.load_state_dict(best_model) - def choose_netG(self): + def choose_netG(self, verbose=False): # compute accept probabilities - print("Networks:", ' '.join(['%3d/%3d' % (A, S) for A, S in zip(self.netGs_accepts, self.netGs_samples)])) - # give initial acceptance probabilities of 1% - probs = np.array([(A+1)/(S+1) for A, S in zip(self.netGs_accepts, self.netGs_samples)]) + # give initially equal acceptance probabilities + if verbose: + print("Networks:", ' '.join(['%3d/%3d' % (A, S) for A, S in zip(self.netGs_accepts, self.netGs_samples)])) + probs = np.array([(A + 1.) / (S + 1) for A, S in zip(self.netGs_accepts, self.netGs_samples)]) probs /= probs.sum() i = np.random.choice(np.arange(len(probs)), p=probs) self.netG.load_state_dict(self.netGs[i]) @@ -227,11 +230,11 @@ def sample( alllatent = [] alllikes = [] allncall = 0 - nsamples = 20 scales = [alpha * 1.0 for _ in self.netGs_init_x] + nsamples = 200 // len(self.netGs_init_x) while len(allsamples) < mcmc_steps: - neti = self.choose_netG() - init_x = self.netGs_init_x[neti] + neti = self.choose_netG(verbose = True) + #init_x = self.netGs_init_x[neti] scale = scales[neti] samples, likes, latent, scale, ncall, accept = self.subsample( @@ -244,6 +247,7 @@ def sample( self.netGs_accepts[neti] += accept.min() self.netGs_samples[neti] += nsamples self.netGs_init_x[neti] = samples[-1] + init_x = samples[-1] scales[neti] = scale allsamples += samples @@ -267,10 +271,6 @@ def sample( if self.writer is not None: self.writer.add_figure('chain', fig, self.total_iters) - if out_chain is not None: - for ib in range(batch_size): - files[ib].close() - return samples, likes, latent, scale, ncall @@ -424,113 +424,11 @@ def transform(x): return x files[ib].write(" ".join(["%.5E" % vi for vi in v[ib]])) files[ib].write("\n") files[ib].flush() + + if out_chain is not None: + for ib in range(batch_size): + files[ib].close() - return samples, likes, latent, scale, ncall, accept - def _jacobian(self, z): - """ Calculate det d f^{-1} (z)/dz - """ - z.requires_grad_(True) - x, _ = self.netG(z, mode='inverse') - J = torch.stack([torch.autograd.grad(outputs=x[:, i], inputs=z, retain_graph=True, - grad_outputs=torch.ones(z.shape[0]))[0] for i in - range(x.shape[1])]).permute(1, 0, 2) - return torch.stack([torch.log(torch.abs(torch.det(J[i, :, :]))) - for i in range(x.shape[0])]) - - def _train(self, epoch, model, loader, noise=0.0): - - model.train() - train_loss = 0 - - for batch_idx, data in enumerate(loader): - if isinstance(data, list): - if len(data) > 1: - cond_data = data[1].float() - cond_data = cond_data.to(self.device) - else: - cond_data = None - data = data[0] - data = (data + noise * torch.randn_like(data)).to(self.device) - self.optimizer.zero_grad() - loss = -model.log_probs(data, cond_data).mean() - train_loss += loss.item() - loss.backward() - self.optimizer.step() - - for module in model.modules(): - if isinstance(module, BatchNormFlow): - module.momentum = 0 - - with torch.no_grad(): - model(loader.dataset.tensors[0].to(data.device)) - - for module in model.modules(): - if isinstance(module, BatchNormFlow): - module.momentum = 1 - - return train_loss / len(loader.dataset) - - def _validate(self, epoch, model, loader): - - model.eval() - val_loss = 0 - cond_data = None - - for batch_idx, data in enumerate(loader): - if isinstance(data, list): - if len(data) > 1: - cond_data = data[1].float() - cond_data = cond_data.to(self.device) - data = data[0] - data = data.to(self.device) - with torch.no_grad(): - # sum up batch loss - val_loss += -model.log_probs(data, cond_data).sum().item() - - return val_loss / len(loader.dataset) - - def _train_plot(self, model, dataset, plot_grid=True): - - model.eval() - with torch.no_grad(): - x_synth = model.sample(dataset.size).detach().cpu().numpy() - x_data = torch.from_numpy(dataset).float() - z, _ = model(x_data) - z = z.detach().cpu().numpy() - if plot_grid: - grid = [] - for x in np.linspace(np.min(dataset[:, 0]), np.max(dataset[:, 0]), 10): - for y in np.linspace(np.min(dataset[:, 1]), np.max(dataset[:, 1]), 5000): - grid.append([x, y]) - for y in np.linspace(np.min(dataset[:, 1]), np.max(dataset[:, 1]), 10): - for x in np.linspace(np.min(dataset[:, 0]), np.max(dataset[:, 0]), 5000): - grid.append([x, y]) - grid = np.array(grid) - x_grid = torch.from_numpy(grid).float() - z_grid, _ = model(x_grid) - z_grid = z_grid.detach().cpu().numpy() - - if self.writer is not None: - fig, ax = plt.subplots(2, figsize=(5, 10)) - ax[0].scatter(dataset[:, 0], dataset[:, 1], c=dataset[:, 0], s=4) - ax[1].scatter(z[:, 0], z[:, 1], c=dataset[:, 0], s=4) - self.writer.add_figure('latent', fig, self.total_iters) + return samples, likes, latent, scale, ncall, accept - if self.path: - fig, ax = plt.subplots(1, 3, figsize=(10, 4)) - if plot_grid: - ax[0].scatter(grid[:, 0], grid[:, 1], c=grid[:, 0], marker='.', s=1, linewidths=0) - ax[0].scatter(dataset[:, 0], dataset[:, 1], s=4) - ax[0].set_title('Real data') - if plot_grid: - ax[1].scatter(z_grid[:, 0], z_grid[:, 1], c=grid[:, 0], marker='.', s=1, linewidths=0) - ax[1].scatter(z[:, 0], z[:, 1], s=4) - ax[1].set_title('Latent data') - ax[1].set_xlim([-3, 3]) - ax[1].set_ylim([-3, 3]) - ax[2].scatter(x_synth[:, 0], x_synth[:, 1], s=2) - ax[2].set_title('Synthetic data') - plt.tight_layout() - plt.savefig(os.path.join(self.path, 'plots', 'plot_%s.png' % self.total_iters)) - plt.close() diff --git a/nnest/networks.py b/nnest/networks.py index f6e4ff9..2c4d52c 100755 --- a/nnest/networks.py +++ b/nnest/networks.py @@ -183,10 +183,10 @@ class SingleSpeed(nn.Module): def __init__(self, num_inputs, num_hidden, num_blocks, num_layers): super(SingleSpeed, self).__init__() + modules = [] mask = torch.arange(0, num_inputs) % 2 mask = mask.float() - modules = [] - for _ in range(num_blocks): + for i in range(num_blocks): modules += [ CouplingLayer( num_inputs, num_hidden, mask, None, diff --git a/nnest/trainer.py b/nnest/trainer.py index 85230dc..472bebf 100644 --- a/nnest/trainer.py +++ b/nnest/trainer.py @@ -410,7 +410,7 @@ def _validate(self, epoch, model, loader): data = data.to(self.device) with torch.no_grad(): # sum up batch loss - val_loss += -model.log_probs(data, cond_data).mean().item() + val_loss += -model.log_probs(data, cond_data).sum().item() return val_loss / len(loader.dataset) From 224551443803b1fadb5d2cf28f79516ed7582854 Mon Sep 17 00:00:00 2001 From: Johannes Buchner Date: Sat, 4 May 2019 12:19:32 +0200 Subject: [PATCH 3/3] combine emcee with neural network --- examples/mcmc/rosenbrock.py | 2 +- nnest/mcmc.py | 10 ++- nnest/multitrainer.py | 120 ++++++++++++++++++++++++------------ 3 files changed, 89 insertions(+), 43 deletions(-) diff --git a/examples/mcmc/rosenbrock.py b/examples/mcmc/rosenbrock.py index b5c4a97..d37ce7d 100644 --- a/examples/mcmc/rosenbrock.py +++ b/examples/mcmc/rosenbrock.py @@ -27,7 +27,7 @@ def transform(x): sampler = MCMCSampler(args.x_dim, loglike, transform=transform, log_dir=args.log_dir, hidden_dim=args.hidden_dim, num_layers=args.num_layers, num_blocks=args.num_blocks, num_slow=args.num_slow) sampler.run(train_iters=args.train_iters, mcmc_steps=args.mcmc_steps, single_thin=10, - bootstrap_iters=args.burnin_iters, bootstrap_mcmc_steps=5000) + bootstrap_iters=args.burnin_iters, bootstrap_mcmc_steps=5000 + 1000 * args.x_dim) if __name__ == '__main__': diff --git a/nnest/mcmc.py b/nnest/mcmc.py index e61a0fd..24cb860 100644 --- a/nnest/mcmc.py +++ b/nnest/mcmc.py @@ -48,7 +48,7 @@ def _init_samples(self, mcmc_steps=5000, mcmc_batch_size=5, ignore_rows=0.3): logl = self.loglike(v) samples = [] likes = [] - for i in range(mcmc_steps): + for i in range(mcmc_steps // mcmc_batch_size): du = np.random.standard_normal(u.shape) * 0.1 u_prime = u + du v_prime = self.transform(u_prime) @@ -73,11 +73,13 @@ def _init_samples(self, mcmc_steps=5000, mcmc_batch_size=5, ignore_rows=0.3): u = u_prime * m + u * (1 - m) v = v_prime * m + v * (1 - m) logl = logl_prime * mask + logl * (1 - mask) + assert v.shape[1] == self.x_dim, v.shape samples.append(v) likes.append(logl) samples = np.transpose(np.array(samples), axes=[1, 0, 2]) loglikes = -np.transpose(np.array(likes), axes=[1, 0]) weights = np.ones(loglikes.shape) + assert samples.shape[-1] == self.x_dim, samples.shape self._chain_stats(samples) self._save_samples(samples, weights, loglikes) names = ['p%i' % i for i in range(int(self.x_dim))] @@ -109,7 +111,7 @@ def run( bootstrap_batch_size=5, alpha=0, single_thin=1, - ignore_rows=0.0): + ignore_rows=0.3): if alpha == 0.0: alpha = 1 / self.x_dim ** 0.5 @@ -132,7 +134,7 @@ def run( loglike=self.loglike, transform=transform, mcmc_steps=bootstrap_mcmc_steps, alpha=alpha, dynamic=False, show_progress=True, - init_x=next_start) + init_x=next_start, plot=True) next_start = samples[:,-1,:] #next_start = None samples = transform(samples) @@ -144,7 +146,9 @@ def run( ignore_rows=ignore_rows) samples = mc.makeSingleSamples(single_thin=single_thin) + print(samples.shape) samples = samples[:, :self.x_dim] + assert samples.shape[1] == self.x_dim mean = np.mean(samples, axis=0) std = np.std(samples, axis=0) samples = (samples - mean) / std diff --git a/nnest/multitrainer.py b/nnest/multitrainer.py index f0c4a37..9394d90 100644 --- a/nnest/multitrainer.py +++ b/nnest/multitrainer.py @@ -7,7 +7,7 @@ from __future__ import print_function from __future__ import division -import os +import os, sys import time import torch @@ -18,6 +18,8 @@ import scipy.spatial from tqdm import tqdm +import emcee + import matplotlib matplotlib.use('Agg') @@ -113,6 +115,8 @@ def train( save_interval=50, noise=0.0, validation_fraction=0.05): + + assert samples.shape[1] == self.x_dim, samples.shape start_time = time.time() @@ -157,6 +161,7 @@ def train( self.netGs_samples = [] self.netGs_init_x = [] + stored = ' ' for epoch in range(1, max_iters + 1): self.total_iters += 1 @@ -165,10 +170,8 @@ def train( epoch, self.netG, train_loader, noise=training_noise) validation_loss = self._validate(epoch, self.netG, valid_loader) - stored = ' ' - # remember only if: - # 1) much better - # 2) when performance is improving drastically at each epoch, avoid every one of them. + # 1) remember only if much better + # 2) when performance is improving drastically at each epoch, avoid storing every epoch. if validation_loss < best_validation_loss - 1.0 and epoch >= best_validation_epoch + save_interval: best_validation_epoch = epoch best_validation_loss = validation_loss @@ -183,6 +186,7 @@ def train( print( 'Epoch: {} validation loss: {:6.4f} {}'.format( epoch, validation_loss, stored)) + stored = ' ' #if epoch % save_interval == 0: @@ -201,7 +205,8 @@ def choose_netG(self, verbose=False): # compute accept probabilities # give initially equal acceptance probabilities if verbose: - print("Networks:", ' '.join(['%3d/%3d' % (A, S) for A, S in zip(self.netGs_accepts, self.netGs_samples)])) + print("Sampling from networks:", ' '.join(['%4d/%4d' % (A, S) for A, S in zip(self.netGs_accepts, self.netGs_samples)]), end=" \r") + sys.stdout.flush() probs = np.array([(A + 1.) / (S + 1) for A, S in zip(self.netGs_accepts, self.netGs_samples)]) probs /= probs.sum() i = np.random.choice(np.arange(len(probs)), p=probs) @@ -224,54 +229,91 @@ def sample( plot=False, out_chain=None, max_prior=None, - max_start_tries=100): + nwalkers=40): + def transformed_logpost(z): + assert z.shape == (self.x_dim,), z.shape + x, log_det_J = self.netG(torch.from_numpy(z.reshape((1,-1))).float().to(self.device), mode='inverse') + x = x.detach().cpu().numpy() + lnlike = float(loglike(transform(x))) + log_det_J = float(log_det_J.detach()) + logprob = log_det_J + lnlike + #print('Like=%.1f' % logprob, x) + return logprob + allsamples = [] alllatent = [] alllikes = [] allncall = 0 - scales = [alpha * 1.0 for _ in self.netGs_init_x] - nsamples = 200 // len(self.netGs_init_x) - while len(allsamples) < mcmc_steps: + populations = [None for net in self.netGs] + nsteps = 4 + #nsamples = 200 // len(self.netGs_init_x) + while allncall < mcmc_steps: neti = self.choose_netG(verbose = True) - #init_x = self.netGs_init_x[neti] - scale = scales[neti] - - samples, likes, latent, scale, ncall, accept = self.subsample( - mcmc_steps=nsamples, alpha=scale, dynamic=dynamic, - loglike=loglike, init_x=init_x, logl=logl, loglstar=loglstar, - transform=transform, show_progress=False, plot=plot, - out_chain=out_chain, max_prior=max_prior, max_start_tries=max_start_tries, - ) - - self.netGs_accepts[neti] += accept.min() - self.netGs_samples[neti] += nsamples - self.netGs_init_x[neti] = samples[-1] - init_x = samples[-1] - scales[neti] = scale + sampler = populations[neti] + if sampler is None: + z0 = np.random.normal(size=(nwalkers, self.x_dim)) + if init_x is not None: + batch_size = init_x.shape[0] + z, _ = self.netG(torch.from_numpy(init_x).float().to(self.device)) + z = z.detach() + z0[:batch_size,:] = z + + sampler = emcee.EnsembleSampler(nwalkers=nwalkers, dim=self.x_dim, lnpostfn=transformed_logpost) + pos, lnprob, rstate = sampler.run_mcmc(z0, nsteps) + populations[neti] = sampler + else: + # advance population + pos, lnprob, rstate = sampler.run_mcmc(sampler.chain[:,-1,:], lnprob0=sampler.lnprobability[:,-1], N=nsteps) - allsamples += samples - alllikes += likes - alllatent += latent - allncall = allncall + ncall - self.logger.info('scales: %s' % scales) - + alllatent.append(pos) + alllikes.append(lnprob) + allncall += nsteps * nwalkers + x, log_det_J = self.netG(torch.from_numpy(pos).float().to(self.device), mode='inverse') + x = x.detach().cpu().numpy() + allsamples.append(x) + + ar = sampler.acceptance_fraction + Nsamples = len(sampler.flatchain) + #self.logger.info('Network %d sampler accepted: %.3f (%.d/%d)' % (neti, ar.mean(), ar.mean() * Nsamples, Nsamples)) + self.netGs_accepts[neti] = int(ar.mean() * Nsamples) + self.netGs_samples[neti] = Nsamples + + #print(np.shape(allsamples), np.shape(alllikes), np.shape(alllatent)) # Transpose so shape is (chain_num, iteration, dim) - samples = np.transpose(np.array(allsamples), axes=[1, 0, 2]) - likes = np.transpose(np.array(alllikes), axes=[1, 0]) - latent = np.transpose(np.array(alllatent), axes=[1, 0, 2]) + #samples = np.transpose(np.array(allsamples), axes=[1, 0, 2]) + #latent = np.transpose(np.array(alllatent), axes=[1, 0, 2]) + #likes = np.transpose(np.array(alllikes), axes=[1, 0]) + samples = np.array(allsamples).reshape((1, -1, self.x_dim)) + latent = np.array(alllatent).reshape((1, -1, self.x_dim)) + likes = np.array(alllikes).reshape((1, -1)) + print() ncall = allncall if self.path and plot: cmap = plt.cm.jet cmap.set_under('w', 1) fig, ax = plt.subplots() - ax.hist2d(samples[:, -1, 0], samples[:, -1, 1], + ax.hist2d(samples[0, :, 0], samples[0, :, 1], bins=200, cmap=cmap, vmin=1, alpha=0.2) - if self.writer is not None: - self.writer.add_figure('chain', fig, self.total_iters) - - return samples, likes, latent, scale, ncall + #if self.writer is not None: + # self.writer.add_figure('chain', fig, self.total_iters) + plt.tight_layout() + plt.savefig(os.path.join(self.path, 'plots', 'chain_%s.png' % self.total_iters)) + plt.close() + fig, ax = plt.subplots() + ax.plot(likes[0, len(likes[0])//3:]) + self.logger.info('lnLike: %d+-%d -> %d+-%d' % ( + alllikes[len(alllikes)//3].mean(), alllikes[len(alllikes)//3].std(), + alllikes[-1].mean(), alllikes[-1].std() + )) + #if self.writer is not None: + # self.writer.add_figure('likeevol', fig, self.total_iters) + plt.tight_layout() + plt.savefig(os.path.join(self.path, 'plots', 'likeevol_%s.png' % self.total_iters)) + plt.close() + + return samples, likes, latent, alpha, ncall def subsample(