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..d37ce7d 100644 --- a/examples/mcmc/rosenbrock.py +++ b/examples/mcmc/rosenbrock.py @@ -11,16 +11,23 @@ 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 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 + 1000 * args.x_dim) if __name__ == '__main__': @@ -30,6 +37,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..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))] @@ -116,7 +118,9 @@ def run( if self.log: self.logger.info('Alpha [%5.4f]' % (alpha)) - + + allsamples = [] + next_start = None for t in range(bootstrap_iters): if t == 0: @@ -128,26 +132,43 @@ 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, plot=True) + next_start = samples[:,-1,:] + #next_start = None 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]], 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 + 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() 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, + 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 new file mode 100644 index 0000000..9394d90 --- /dev/null +++ b/nnest/multitrainer.py @@ -0,0 +1,476 @@ +""" +.. 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, sys +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 emcee + +import matplotlib + +matplotlib.use('Agg') +import matplotlib.pyplot as plt + +from nnest.networks import SingleSpeed, FastSlow, BatchNormFlow +from nnest.utils.logger import create_logger +from nnest.trainer import Trainer + + +class MultiTrainer(Trainer): + + 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): + + assert samples.shape[1] == self.x_dim, samples.shape + + 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 = 0.5 * np.mean(dists) / self.x_dim + 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 = [] + + stored = ' ' + 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) + + # 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 + 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)) + 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) + ) + if self.x_dim == 2: + self._train_plot(self.netG, samples) + self.netG.load_state_dict(best_model) + + def choose_netG(self, verbose=False): + # compute accept probabilities + # give initially equal acceptance probabilities + if verbose: + 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) + 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, + 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 + 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) + 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) + + 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]) + #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[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) + 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( + 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() + + if out_chain is not None: + for ib in range(batch_size): + files[ib].close() + + + return samples, likes, latent, scale, ncall, accept + 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/sampler.py b/nnest/sampler.py index 19a9443..e2e2abc 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 a5733a9..472bebf 100644 --- a/nnest/trainer.py +++ b/nnest/trainer.py @@ -9,7 +9,6 @@ import os import time -import copy import torch from torch.utils.data import DataLoader @@ -63,44 +62,48 @@ 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) + 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 - if self.path is not None: - self.logger.info(self.netG) - self.writer = SummaryWriter(self.path) - def train( self, samples, @@ -108,7 +111,7 @@ def train( log_interval=50, save_interval=50, noise=0.0, - validation_fraction=0.1): + validation_fraction=0.05): start_time = time.time() @@ -128,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( @@ -146,7 +149,7 @@ def train( best_validation_loss = float('inf') best_validation_epoch = 0 - best_model = copy.deepcopy(self.netG) + best_model = self.netG.state_dict() for epoch in range(1, max_iters + 1): @@ -159,7 +162,7 @@ def train( if validation_loss < best_validation_loss: best_validation_epoch = epoch best_validation_loss = validation_loss - best_model = copy.deepcopy(self.netG) + best_model = self.netG.state_dict() if epoch == 1 or epoch % log_interval == 0: self.logger.info('Epoch [%i] train loss [%5.4f] validation loss [%5.4f]' % ( @@ -175,8 +178,7 @@ def train( self._train_plot(self.netG, samples) self.logger.info('Best epoch [%i] validation loss [%5.4f]' % (best_validation_epoch, best_validation_loss)) - - self.netG.load_state_dict(best_model.state_dict()) + self.netG.load_state_dict(best_model) def sample( self, @@ -408,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)