def main(args): # load data print('loading training data...') dataset_directory = get_data_directory(__file__) dataset_path = os.path.join(dataset_directory, 'faces_training.csv') if not os.path.exists(dataset_path): try: os.makedirs(dataset_directory) except OSError as e: if e.errno != errno.EEXIST: raise pass wget.download('https://d2hg8soec8ck9v.cloudfront.net/datasets/faces_training.csv', dataset_path) data = torch.tensor(np.loadtxt(dataset_path, delimiter=',')).float() sparse_gamma_def = SparseGammaDEF() # Due to the special logic in the custom guide (e.g. parameter clipping), the custom guide # seems to be more amenable to higher learning rates. # Nevertheless, the easy guide performs the best (presumably because of numerical instabilities # related to the gamma distribution in the custom guide). learning_rate = 0.2 if args.guide in ['auto', 'easy'] else 4.5 momentum = 0.05 if args.guide in ['auto', 'easy'] else 0.1 opt = optim.AdagradRMSProp({"eta": learning_rate, "t": momentum}) # use one of our three different guide types if args.guide == 'auto': guide = AutoDiagonalNormal(sparse_gamma_def.model, init_loc_fn=init_to_feasible) elif args.guide == 'easy': guide = MyEasyGuide(sparse_gamma_def.model) else: guide = sparse_gamma_def.guide # this is the svi object we use during training; we use TraceMeanField_ELBO to # get analytic KL divergences svi = SVI(sparse_gamma_def.model, guide, opt, loss=TraceMeanField_ELBO()) # we use svi_eval during evaluation; since we took care to write down our model in # a fully vectorized way, this computation can be done efficiently with large tensor ops svi_eval = SVI(sparse_gamma_def.model, guide, opt, loss=TraceMeanField_ELBO(num_particles=args.eval_particles, vectorize_particles=True)) print('\nbeginning training with %s guide...' % args.guide) # the training loop for k in range(args.num_epochs): loss = svi.step(data) # for the custom guide we clip parameters after each gradient step if args.guide == 'custom': clip_params() if k % args.eval_frequency == 0 and k > 0 or k == args.num_epochs - 1: loss = svi_eval.evaluate_loss(data) print("[epoch %04d] training elbo: %.4g" % (k, -loss))
def main(args): # load data print('loading training data...') dataset_directory = get_data_directory(__file__) dataset_path = os.path.join(dataset_directory, 'faces_training.csv') if not os.path.exists(dataset_path): try: os.makedirs(dataset_directory) except OSError as e: if e.errno != errno.EEXIST: raise pass wget.download('https://d2fefpcigoriu7.cloudfront.net/datasets/faces_training.csv', dataset_path) data = torch.tensor(np.loadtxt(dataset_path, delimiter=',')).float() sparse_gamma_def = SparseGammaDEF() # due to the special logic in the custom guide (e.g. parameter clipping), the custom guide # is more numerically stable and enables us to use a larger learning rate (and consequently # achieves better results) learning_rate = 0.2 if args.auto_guide else 4.5 momentum = 0.05 if args.auto_guide else 0.1 opt = optim.AdagradRMSProp({"eta": learning_rate, "t": momentum}) # either use an automatically constructed guide (see pyro.contrib.autoguide for details) or our custom guide guide = AutoDiagonalNormal(sparse_gamma_def.model) if args.auto_guide else sparse_gamma_def.guide # this is the svi object we use during training; we use TraceMeanField_ELBO to # get analytic KL divergences svi = SVI(sparse_gamma_def.model, guide, opt, loss=TraceMeanField_ELBO()) # we use svi_eval during evaluation; since we took care to write down our model in # a fully vectorized way, this computation can be done efficiently with large tensor ops svi_eval = SVI(sparse_gamma_def.model, guide, opt, loss=TraceMeanField_ELBO(num_particles=args.eval_particles, vectorize_particles=True)) guide_description = 'automatically constructed' if args.auto_guide else 'custom' print('\nbeginning training with %s guide...' % guide_description) # the training loop for k in range(args.num_epochs): loss = svi.step(data) if not args.auto_guide: # for the custom guide we clip parameters after each gradient step sparse_gamma_def.clip_params() if k % args.eval_frequency == 0 and k > 0 or k == args.num_epochs - 1: loss = svi_eval.evaluate_loss(data) print("[epoch %04d] training elbo: %.4g" % (k, -loss))
def getELBO(self, elbo_type, particles): if elbo_type == 'TraceELBO': return Trace_ELBO(num_particles=particles) elif elbo_type == "MeanFieldELBO": return TraceMeanField_ELBO(num_particles=particles) else: raise ValueError("{} ELBO not supported".format(elbo_type))
def train(gpmodule, optimizer=None, loss_fn=None, retain_graph=None, num_steps=1000): """ A helper to optimize parameters for a GP module. :param ~pyro.contrib.gp.models.GPModel gpmodule: A GP module. :param ~torch.optim.Optimizer optimizer: A PyTorch optimizer instance. By default, we use Adam with ``lr=0.01``. :param callable loss_fn: A loss function which takes inputs are ``gpmodule.model``, ``gpmodule.guide``, and returns ELBO loss. By default, ``loss_fn=TraceMeanField_ELBO().differentiable_loss``. :param bool retain_graph: An optional flag of ``torch.autograd.backward``. :param int num_steps: Number of steps to run SVI. :returns: a list of losses during the training procedure :rtype: list """ optimizer = ( torch.optim.Adam(gpmodule.parameters(), lr=0.01) if optimizer is None else optimizer ) # TODO: add support for JIT loss loss_fn = TraceMeanField_ELBO().differentiable_loss if loss_fn is None else loss_fn def closure(): optimizer.zero_grad() loss = loss_fn(gpmodule.model, gpmodule.guide) torch_backward(loss, retain_graph) return loss losses = [] for i in range(num_steps): loss = optimizer.step(closure) losses.append(torch_item(loss)) return losses
def main(args): # load data print('loading training data...') dataset_directory = get_data_directory(__file__) dataset_path = os.path.join(dataset_directory, 'faces_training.csv') if not os.path.exists(dataset_path): try: os.makedirs(dataset_directory) except OSError as e: if e.errno != errno.EEXIST: raise pass wget.download( 'https://d2fefpcigoriu7.cloudfront.net/datasets/faces_training.csv', dataset_path) data = torch.tensor(np.loadtxt(dataset_path, delimiter=',')).float() learning_rate = 4.5 momentum = 0.1 opt = optim.AdagradRMSProp({"eta": learning_rate, "t": momentum}) # this is the svi object we use during training; we use TraceMeanField_ELBO to # get analytic KL divergences svi = SVI(model, guide, opt, loss=TraceMeanField_ELBO()) # we use svi_eval during evaluation; since we took care to write down our model in # a fully vectorized way, this computation can be done efficiently with large tensor ops svi_eval = SVI(model_original, guide, opt, loss=TraceMeanField_ELBO(num_particles=args.eval_particles, vectorize_particles=True)) guide_description = 'custom' print('\nbeginning training with %s guide...' % guide_description) # the training loop for k in range(args.num_epochs): loss = svi.step(data) clip_params() if k % args.eval_frequency == 0 and k > 0 or k == args.num_epochs - 1: loss = svi_eval.evaluate_loss(data) print("[epoch %04d] training elbo: %.4g" % (k, -loss))
def __init__(self, net, prior, guide_builder=None, name="", closed_form_kl=True): super().__init__(net, prior, guide_builder=guide_builder, name=name) self.cached_output = None self.cached_kl_loss = None self._loss = TraceMeanField_ELBO() if closed_form_kl else Trace_ELBO()
def _train_svi(self, train_loader, epochs, lr, device): self.device = device print("\n == SVI training ==") optimizer = pyro.optim.Adam({"lr": lr}) elbo = TraceMeanField_ELBO() svi = SVI(self.model, self.guide, optimizer, loss=elbo) loss_list = [] accuracy_list = [] start = time.time() for epoch in range(epochs): loss = 0.0 correct_predictions = 0.0 for x_batch, y_batch in train_loader: x_batch = x_batch.to(device) y_batch = y_batch.to(device) labels = y_batch.argmax(-1) loss += svi.step(x_data=x_batch, y_data=labels) outputs = self.forward(x_batch, n_samples=10) predictions = outputs.argmax(dim=-1) correct_predictions += (predictions == labels).sum().item() if DEBUG: print("\n", pyro.get_param_store()["model.0.weight_loc"][0][:5]) print("\n", predictions[:10], "\n", labels[:10]) total_loss = loss / len(train_loader.dataset) accuracy = 100 * correct_predictions / len(train_loader.dataset) print( f"\n[Epoch {epoch + 1}]\t loss: {total_loss:.2f} \t accuracy: {accuracy:.2f}", end="\t") loss_list.append(loss) accuracy_list.append(accuracy) execution_time(start=start, end=time.time()) self.save() plot_loss_accuracy(dict={ 'loss': loss_list, 'accuracy': accuracy_list }, path=TESTS + self.name + "/" + self.name + "_training.png")
def train(self, *, raw_expr, encoded_expr, num_epochs=100, batch_size=32, learning_rate=1e-3, posterior_samples=20): logging.info('Validating data ...') assert (raw_expr.shape == encoded_expr.shape) read_depth = torch.tensor(raw_expr.sum(-1)[:, np.newaxis]).to(self.device) raw_expr = torch.tensor(raw_expr).to(self.device) encoded_expr = torch.tensor(encoded_expr).to(self.device) logging.info('Initializing model ...') pyro.clear_param_store() adam_params = {"lr": 1e-3} optimizer = Adam(adam_params) svi = SVI(self.model, self.guide, optimizer, loss=TraceMeanField_ELBO()) logging.info('Training ...') num_batches = read_depth.shape[0] // batch_size bar = trange(num_epochs) try: for epoch in bar: running_loss = 0.0 for batch in self.epoch_batch(raw_expr, encoded_expr, read_depth, batch_size=batch_size): loss = svi.step(*batch) running_loss += loss / batch_size bar.set_postfix(epoch_loss='{:.2e}'.format(running_loss)) except KeyboardInterrupt: logging.error('Interrupted training.') self.summarize_posterior(raw_expr, encoded_expr, read_depth, num_samples=posterior_samples) return self
def log_prob(self, window_begin, window_end, truth): forecast_hours = len(truth) self.args.funsor = False # sets behavior of model and guide assert 0 <= window_begin < window_end < window_end + forecast_hours <= len( self.counts) features = self.features[window_begin: window_end + forecast_hours] \ .to(device=self.args.device) x = self.counts[window_begin:window_end].to(device=self.args.device) y = truth.to(device=self.args.device) xy = torch.cat([x, y], dim=0) loss = TraceMeanField_ELBO().loss logp_x = -loss(self.model, self.guide, features[:len(x)], x) logp_xy = -loss(self.model, self.guide, features[:len(xy)], xy) return logp_xy - logp_x
def train(self, *, accessibility_matrix, num_epochs=125, batch_size=32, learning_rate=1e-3, eval_every=10, test_proportion=0.05): logging.info('Initializing model ...') pyro.clear_param_store() optimizer = Adam({"lr": 1e-3}) self.svi = SVI(self.model, self.guide, optimizer, loss=TraceMeanField_ELBO()) test_set = np.random.rand( accessibility_matrix.shape[0]) < test_proportion train_set = ~test_set logging.info("Training with {} cells, testing with {}.".format( str(train_set.sum()), str(test_set.sum()))) num_batches = accessibility_matrix.shape[0] // batch_size logging.info('Training for {} epochs'.format(str(num_epochs))) try: for epoch in range(1, num_epochs + 1): running_loss = 0.0 for batch in self.epoch_batch(accessibility_matrix[train_set], batch_size=batch_size): loss = self.svi.step(*batch) running_loss += loss / batch_size logging.info('Done epoch {}/{}. Training loss: {:.3e}'.format( str(epoch), str(num_epochs), running_loss)) if (epoch % eval_every == 0 or epoch == num_epochs) and test_set.sum() > 0: test_logp = self.evaluate(accessibility_matrix[test_set], batch_size=batch_size) logging.info('Test logp: {:.4e}'.format(test_logp)) except KeyboardInterrupt: logging.error('Interrupted training.') self.summarize_posterior(accessibility_matrix) return self
def fit(self, data_loader, optim, num_epochs, callback=None, num_particles=1, closed_form_kl=True, device=None): """Optimizes the variational parameters on data from data_loader using optim for num_epochs. :param Iterable data_loader: iterable over batches of data, e.g. a torch.utils.data.DataLoader. Assumes that each element consists of a length two tuple of list, with the first element either containing a single object or a list of objects, e.g. torch.Tensors, that are the inputs to the neural network. The second element is a single torch.Tensor e.g. of class labels. :param optim: pyro optimizer to be used for constructing an SVI object, e.g. pyro.optim.Adam({"lr": 1e-3}). :param int num_epochs: number of passes over data_loader. :param callable callback: optional function to invoke after every training epoch. Receives the BNN object, the epoch number and the average value of the ELBO over the epoch. May return True to terminate optimization before num_epochs, e.g. if it finds that a validation log likelihood saturates. :param int num_particles: number of MC samples for estimating the ELBO. :param bool closed_form_kl: whether to use TraceMeanField_ELBO or Trace_ELBO, i.e. calculate KL divergence between approximate posterior and prior in closed form or via a Monte Carlo estimate. :param torch.device device: optional device to send the data to. """ old_training_state = self.net.training self.net.train(True) loss = TraceMeanField_ELBO( num_particles) if closed_form_kl else Trace_ELBO(num_particles) svi = SVI(self.model, self.guide, optim, loss=loss) for i in range(num_epochs): elbo = 0. num_batch = 1 for num_batch, (input_data, observation_data) in enumerate( iter(data_loader), 1): elbo += svi.step(tuple(_to(input_data, device)), tuple(_to(observation_data, device))[0]) # the callback can stop training by returning True if callback is not None and callback(self, i, elbo / num_batch): break self.net.train(old_training_state) return svi
def train(self, data_loader, n_epochs, num_particles=1, lr=1e-3, log_per=5, show_smooth=True, save_per=10): pyro.clear_param_store() self.guide = partial(guide, model=self) svi = SVI(self, self.guide, Adam({"lr": lr}), TraceMeanField_ELBO(num_particles=num_particles)) losses = [] fig = None pp = ProgressPlotter(losses, "$-ELBO$", log_per, show_smooth) pp.start() for epoch in range(n_epochs): total_loss = 0. for x, y in data_loader: x, y = x.float().to(self.device), y.float().to(self.device) loss = svi.step(x, y) / y.numel() total_loss += loss total_loss /= len(data_loader) losses.append(total_loss) fig = pp.update(epoch) if epoch % save_per == 1: self.save('MultilayerBayesianTrick.pth') return pp.fig, losses
def infer_parameters(self, loader, lr=0.01, momentum=0.9, num_epochs=30): optim = pyroopt.SGD({'lr': lr, 'momentum': momentum, 'nesterov': True}) elbo = TraceMeanField_ELBO() svi = SVI(self.model, self.guide, optim, elbo) kl_factor = loader.batch_size / len(loader.dataset) for i in range(num_epochs): total_loss = 0.0 total = 0.0 correct = 0.0 for images, labels in loader: loss = svi.step(images.cuda(), labels.cuda(), kl_factor=kl_factor) pred = self.forward(images.cuda(), n_samples=1).mean(0) total_loss += loss / len(loader.dataset) total += labels.size(0) correct += (pred.argmax(-1) == labels.cuda()).sum().item() param_store = pyro.get_param_store() print( f"[Epoch {i + 1}] loss: {total_loss:.5E} accuracy: {correct / total * 100:.5f}" )
def train(self, *, raw_expr, encoded_expr, num_epochs=100, batch_size=32, learning_rate=1e-3, eval_every=10, test_proportion=0.05, use_l1=False, l1_lam=0): seed = 2556 torch.manual_seed(seed) pyro.set_rng_seed(seed) pyro.clear_param_store() logging.info('Validating data ...') assert (raw_expr.shape == encoded_expr.shape) read_depth = raw_expr.sum(-1)[:, np.newaxis] encoded_expr = np.hstack([encoded_expr, np.log(read_depth)]) read_depth = torch.tensor(read_depth).to(self.device) raw_expr = torch.tensor(raw_expr).to(self.device) encoded_expr = torch.tensor(encoded_expr).to(self.device) logging.info('Initializing model ...') self.optimizer = Adam({"lr": 1e-3}) self.loss = TraceMeanField_ELBO() if not use_l1: logging.info('No L1 regularization.') svi = SVI(self.model, self.guide, self.optimizer, loss=self.loss) test_set = np.random.rand(read_depth.shape[0]) < test_proportion train_set = ~test_set logging.info("Training with {} cells, testing with {}.".format( str(train_set.sum()), str(test_set.sum()))) logging.info('Training ...') try: for epoch in range(1, num_epochs + 1): running_loss = 0.0 for batch in self.epoch_batch(raw_expr[train_set], encoded_expr[train_set], read_depth[train_set], batch_size=batch_size): if use_l1: loss = self.custom_step(*batch) else: loss = svi.step(*batch) running_loss += loss / batch_size logging.info('Done epoch {}/{}. Training loss: {:.3e}'.format( str(epoch), str(num_epochs), running_loss)) if (epoch % eval_every == 0 or epoch == num_epochs) and test_set.sum() > 0: test_logp = self.evaluate(raw_expr[test_set], encoded_expr[test_set], read_depth[test_set]) logging.info('Test logp: {:.4e}'.format(test_logp)) except KeyboardInterrupt: logging.error('Interrupted training.') self.summarize_posterior(raw_expr, encoded_expr, read_depth) return self
Evaluate DMEGP on PHYSIONET test dataset. """ # system configuration torch.manual_seed(1) # data configuration dataloaders = fetch_dataloaders_PhysioNet(['val', 'test'], 'HR') dataloader_val = dataloaders['val'] dataloader_test = dataloaders['test'] # load a model model_path = './experiments/physionet/dmegp_physionet.pth' cpt = torch.load(model_path) embed_fn = MLP_embed(cpt['input_dim'], cpt['feature_dim']) mean_fn = MLP(cpt['feature_dim'], cpt['hidden_dim'], cpt['output_dim']) model = DMEGP(cpt['input_dim'], feature_dim=cpt['feature_dim'], mean_fn=mean_fn, embed_fn=embed_fn) model.load_state_dict(cpt['state_dict']) elbo = TraceMeanField_ELBO() loss_fn = elbo.differentiable_loss # evaluate the model rmse_val, loss_val = evaluate(model, dataloader_val, loss_fn, cpt['n_adapt'], cpt['inner_lr']) rmse_te, loss_te = evaluate(model, dataloader_test, loss_fn, cpt['n_adapt'], cpt['inner_lr']) print('[Val] RMSE={:05.3f} and loss={:05.3f}'.format(rmse_val, loss_val)) print('[Test] RMSE={:05.3f} and loss={:05.3f}'.format(rmse_te, loss_te))
def train(args, dataset): """ Train a model and guide to fit a dataset. """ counts = dataset["counts"] num_stations = len(dataset["stations"]) train_size = args.truncate if args.truncate else len(counts) logging.info( "Training on {} stations over {}/{} hours, {} batches/epoch".format( num_stations, train_size, len(counts), int(math.ceil(train_size / args.batch_size)))) time_features = make_time_features(args, 0, len(counts)) control_features = (counts.max(1)[0] + counts.max(2)[0]).clamp(max=1) logging.debug( "On average {:0.1f}/{} stations are open at any one time".format( control_features.sum(-1).mean(), num_stations)) features = torch.cat([time_features, control_features], -1) feature_dim = features.size(-1) logging.debug("feature_dim = {}".format(feature_dim)) metadata = {"args": args, "losses": [], "control": control_features} torch.save(metadata, args.training_filename) if args.device.startswith("cuda"): torch.set_default_tensor_type('torch.cuda.FloatTensor') def optim_config(module_name, param_name): config = { "lr": args.learning_rate, "betas": (0.8, 0.99), "lrd": 0.1**(1 / args.num_steps), } if param_name == "init_scale": config["lr"] *= 0.1 # init_dist sees much less data per minibatch return config training_counts = counts[:args.truncate] if args.truncate else counts data_size = len(training_counts) model = Model(args, features, training_counts).to(device=args.device) guide = Guide(args, features, training_counts).to(device=args.device) optimizer = ClippedAdam(optim_config) if args.funsor: elbo = Funsor_ELBO(args) elif args.analytic_kl: elbo = TraceMeanField_ELBO() else: elbo = Trace_ELBO() svi = SVI(model, guide, optimizer, elbo) losses = [] forecaster = None for step in range(args.num_steps): begin_time = torch.randint(max(1, data_size - args.batch_size), ()).item() end_time = min(data_size, begin_time + args.batch_size) feature_batch = features[begin_time:end_time].to(device=args.device) counts_batch = counts[begin_time:end_time].to(device=args.device) loss = svi.step(feature_batch, counts_batch) / counts_batch.numel() assert math.isfinite(loss), loss losses.append(loss) logging.debug("step {} loss = {:0.4g}".format(step, loss)) if step % 20 == 0: # Save state every few steps. pyro.get_param_store().save(args.param_store_filename) metadata = { "args": args, "losses": losses, "control": control_features } torch.save(metadata, args.training_filename) forecaster = Forecaster(args, dataset, features, model, guide) torch.save(forecaster, args.forecaster_filename) if logging.Logger(None).isEnabledFor(logging.DEBUG): init_scale = pyro.param("init_scale").data trans_scale = pyro.param("trans_scale").data trans_matrix = pyro.param("trans_matrix").data eigs = trans_matrix.eig()[0].norm(dim=-1).sort( descending=True).values logging.debug("guide.diag_part = {}".format( guide.diag_part.data.squeeze())) logging.debug( "init scale min/mean/max: {:0.3g} {:0.3g} {:0.3g}".format( init_scale.min(), init_scale.mean(), init_scale.max())) logging.debug( "trans scale min/mean/max: {:0.3g} {:0.3g} {:0.3g}".format( trans_scale.min(), trans_scale.mean(), trans_scale.max())) logging.debug("trans mat eig:\n{}".format(eigs)) return forecaster
def test_elbo_analytic_kl(self): self.do_elbo_test(True, 3000, TraceMeanField_ELBO())
article_nav_bows = np.zeros( (len(training_data.articles), len(training_data.nav_vocab)), dtype=np.float32) for article_id in range(len(training_data.article_navs)): for nav_id in training_data.article_navs[article_id]: article_nav_bows[article_id][nav_id] += 1 navLDA = NavEmbeddingLDA(device, nav_topics, float(args.dropout), training_data).to(device) optimizer = pyro.optim.Adam({"lr": float(args.lr)}) svi = SVI(navLDA.model, navLDA.guide, optimizer, loss=TraceMeanField_ELBO()) elbos = [] num_batches = int(article_nav_bows.shape[0] / batch_size) for epoch in range(int(args.num_epochs)): article_ids = np.array((range(article_nav_bows.shape[0]))) np.random.shuffle(article_ids) epoch_elbo = 0.0 for batch_id in range(num_batches): batch_article_ids = article_ids[list( range(batch_id * batch_size, (batch_id + 1) * batch_size))] batch = torch.tensor( article_nav_bows[batch_article_ids]).to(device) elbos.append( svi.step(batch, training_data.nav_embeddings,
def test_constant_kl(): model = lambda: pyro.sample("a", dist.Normal(0, 1.)) guide = tyxe.guides.AutoNormal(model) elbo = TraceMeanField_ELBO() assert elbo.loss(model, guide) == elbo.loss(model, guide)
def main(): assert pyro.__version__.startswith('1.6.0') # Enable smoke test to test functionality smoke_test = False logging.info(f"CUDA available: {torch.cuda.is_available()}") # Loading data logging.info("Loading data...") docs = prepro_file_load("doc_word_matrix").to_dense() id2word = prepro_file_load("id2word") # Put vocab into dataframe for exploration of data vocab = pd.DataFrame(columns=['index', 'word']) vocab['index'] = list(id2word.keys()) vocab['word'] = list(id2word.values()) logging.info(f"Dictionary size: {len(vocab)}") logging.info(f"Corpus size: {docs.shape}") # Setting global variables seed = 0 torch.manual_seed(seed) pyro.set_rng_seed(seed) device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") docs = docs.float() num_categories = 0 num_topics = 30 if not smoke_test else 3 batch_size = 32 learning_rate = 1e-3 num_epochs = 50 if not smoke_test else 1 # Training pyro.clear_param_store() prodLDA = ProdLDA(vocab_size=docs.shape[1], num_topics=num_topics, hidden=100 if not smoke_test else 10, dropout=0.2) prodLDA.to(device) optimizer = pyro.optim.Adam({"lr": learning_rate}) svi = SVI(prodLDA.model, prodLDA.guide, optimizer, loss=TraceMeanField_ELBO()) num_batches = int(math.ceil(docs.shape[0] / batch_size)) if not smoke_test else 1 losses = [] logging.info("Training...") bar = trange(num_epochs) for epoch in bar: running_loss = 0.0 for i in range(num_batches): batch_docs = docs[i * batch_size:(i + 1) * batch_size, :].to(device) loss = svi.step(batch_docs) running_loss += loss / batch_docs.size(0) # Save and log losses losses.append(running_loss) bar.set_postfix(epoch_loss='{:.2e}'.format(running_loss)) if epoch % 5 == 0: logging.info('{: >5d}\t{}'.format(epoch, '{:.2e}'.format(running_loss))) logging.info(f"Final loss: {'{:.2e}'.format(losses[-1])}/{losses[-1]}") if not smoke_test: # Plot loss over epochs plt.plot(losses) plt.title("ELBO") plt.xlabel("Epoch") plt.ylabel("Loss") plot_file_name = "../ProdLDA-loss-2017_categories-" + str(num_categories) + \ "_topics-" + str(num_topics) + \ "_batch-" + str(batch_size) + \ "_lr-" + str(learning_rate) + \ "_epochs-" + str(num_epochs) + \ ".png" plt.savefig(plot_file_name) plt.show() # Logging top 10 weighted words in topics beta = prodLDA.beta() for n in range(beta.shape[0]): sorted_, indices = torch.sort(beta[n], descending=True) df = pd.DataFrame(indices[:10].numpy(), columns=['index']) words = pd.merge(df, vocab[['index', 'word']], how='left', on='index')['word'].values.tolist() logging.info(f"Topic {n}: {words}")