def _build_svi(self, loss=None): def per_param_callable(module_name, param_name): params = { 'eps': 1e-5, 'amsgrad': self.hparams.use_amsgrad, 'weight_decay': self.hparams.l2 } if module_name == 'intensity_flow_components' or module_name == 'thickness_flow_components': params['lr'] = self.hparams.pgm_lr return params else: params['lr'] = self.hparams.lr return params if loss is None: loss = self.svi_loss if self.hparams.use_cf_guide: def guide(*args, **kwargs): return self.pyro_model.counterfactual_guide( *args, **kwargs, counterfactual_type=self.hparams.cf_elbo_type) self.svi = SVI(self.pyro_model.svi_model, guide, Adam(per_param_callable), loss) else: self.svi = SVI(self.pyro_model.svi_model, self.pyro_model.svi_guide, Adam(per_param_callable), loss) self.svi.loss_class = loss
def fit_GP(self): ### train every GPf and GPo, cache necessary variables for further filtering self.GPs_losses = [] self.GPo_losses = [] # TODO CACHE DIFFERENT STUFF IF USE SPARSE GP pyro.clear_param_store() num_steps = 2500 for (i, GPs) in enumerate(self.state_transition_model_list): losses = GPs.optimize(optimizer=Adam({"lr": 0.005}), num_steps=num_steps) self.GPs_losses.append(losses) print("training for state transition model {} is done!".format(i)) for (i ,GPo) in enumerate(self.observation_model_list): losses = GPo.optimize(optimizer=Adam({"lr": 0.005}), num_steps=num_steps) self.GPo_losses.append(losses) print("training for observation model {} is done!".format(i)) self.cache_variable() # save the mode self.save_model() return self.GPs_losses, self.GPo_losses
def _build_svi(self, loss=None): def per_param_callable(module_name, param_name): params = { 'eps': 1e-5, 'amsgrad': self.hparams.use_amsgrad, 'weight_decay': self.hparams.l2 } if 'flow_components' in module_name or 'sex_logits' in param_name: params['lr'] = self.hparams.pgm_lr else: params['lr'] = self.hparams.lr print( f'building opt for {module_name} - {param_name} with p: {params}' ) return params if loss is None: loss = self.svi_loss if self.hparams.use_cf_guide: def guide(*args, **kwargs): return self.pyro_model.counterfactual_guide( *args, **kwargs, counterfactual_type=self.hparams.cf_elbo_type) self.svi = SVI(self.pyro_model.svi_model, guide, Adam(per_param_callable), loss) else: self.svi = SVI(self.pyro_model.svi_model, self.pyro_model.svi_guide, Adam(per_param_callable), loss) self.svi.loss_class = loss
def test_beta_bernoulli(Elbo, vectorized): pyro.clear_param_store() data = torch.tensor([1.0] * 6 + [0.0] * 4) def model1(data): alpha0 = torch.tensor(10.0) beta0 = torch.tensor(10.0) f = pyro.sample("latent_fairness", dist.Beta(alpha0, beta0)) for i in pyro.irange("irange", len(data)): pyro.sample("obs_{}".format(i), dist.Bernoulli(f), obs=data[i]) def model2(data): alpha0 = torch.tensor(10.0) beta0 = torch.tensor(10.0) f = pyro.sample("latent_fairness", dist.Beta(alpha0, beta0)) pyro.sample("obs", dist.Bernoulli(f).expand_by(data.shape).independent(1), obs=data) model = model2 if vectorized else model1 def guide(data): alpha_q = pyro.param("alpha_q", torch.tensor(15.0), constraint=constraints.positive) beta_q = pyro.param("beta_q", torch.tensor(15.0), constraint=constraints.positive) pyro.sample("latent_fairness", dist.Beta(alpha_q, beta_q)) elbo = Elbo(num_particles=7, strict_enumeration_warning=False) optim = Adam({"lr": 0.0005, "betas": (0.90, 0.999)}) svi = SVI(model, guide, optim, elbo) for step in range(40): svi.step(data)
def test_svi_multi(): args = make_args() args.assignment_grad = True detections = generate_data(args) pyro.clear_param_store() pyro.param('noise_scale', torch.tensor(args.init_noise_scale), constraint=constraints.positive) pyro.param('objects_loc', torch.randn(args.max_num_objects, 1)) # Learn object_loc via Newton and noise_scale via Adam. elbo = TraceEnum_ELBO(max_plate_nesting=2) adam = Adam({'lr': 0.1}) newton = Newton(trust_radii={'objects_loc': 1.0}) optim = MixedMultiOptimizer([(['noise_scale'], adam), (['objects_loc'], newton)]) for svi_step in range(50): with poutine.trace(param_only=True) as param_capture: loss = elbo.differentiable_loss(model, guide, detections, args) params = { name: pyro.param(name).unconstrained() for name in param_capture.trace.nodes.keys() } optim.step(loss, params) logger.debug( 'step {: >2d}, loss = {:0.6f}, noise_scale = {:0.6f}'.format( svi_step, loss.item(), pyro.param('noise_scale').item()))
def test_svi(): n = 11 data = torch.zeros(n) def model(data): loc = pyro.param('loc', torch.zeros(n, requires_grad=True)) scale = pyro.param('log_scale', torch.zeros(n, requires_grad=True)).exp() pyro.sample('obs', dist.Normal(loc, scale).independent(1), obs=data) def guide(data): pass optim = Adam({'lr': 1e-3}) inference = SVI(model, guide, optim, Trace_ELBO()) counts = [] gc.collect() gc.collect() expected = count_objects_of_type(Trace) for _ in range(10): inference.step(data) counts.append(count_objects_of_type(Trace)) assert set(counts) == set([expected]), counts
def test_subsample_model(auto_class): def model(x, y=None, batch_size=None): loc = pyro.param("loc", lambda: torch.tensor(0.)) scale = pyro.param("scale", lambda: torch.tensor(1.), constraint=constraints.positive) with pyro.plate("batch", len(x), subsample_size=batch_size): batch_x = pyro.subsample(x, event_dim=0) batch_y = pyro.subsample(y, event_dim=0) if y is not None else None mean = loc + scale * batch_x sigma = pyro.sample("sigma", dist.LogNormal(0., 1.)) return pyro.sample("obs", dist.Normal(mean, sigma), obs=batch_y) guide = auto_class(model) full_size = 50 batch_size = 20 pyro.set_rng_seed(123456789) x = torch.randn(full_size) with torch.no_grad(): y = model(x) assert y.shape == x.shape pyro.get_param_store().clear() pyro.set_rng_seed(123456789) svi = SVI(model, guide, Adam({"lr": 0.02}), Trace_ELBO()) for step in range(5): svi.step(x, y, batch_size=batch_size)
def test_median(auto_class, Elbo): def model(): pyro.sample("x", dist.Normal(0.0, 1.0)) pyro.sample("y", dist.LogNormal(0.0, 1.0)) pyro.sample("z", dist.Beta(2.0, 2.0)) guide = auto_class(model) optim = Adam({'lr': 0.05, 'betas': (0.8, 0.99)}) elbo = Elbo(strict_enumeration_warning=False, num_particles=100, vectorize_particles=True) infer = SVI(model, guide, optim, elbo) for _ in range(100): infer.step() if auto_class is AutoLaplaceApproximation: guide = guide.laplace_approximation() median = guide.median() assert_equal(median["x"], torch.tensor(0.0), prec=0.1) if auto_class is AutoDelta: assert_equal(median["y"], torch.tensor(-1.0).exp(), prec=0.1) else: assert_equal(median["y"], torch.tensor(1.0), prec=0.1) assert_equal(median["z"], torch.tensor(0.5), prec=0.1)
def training(): optim = Adam({"lr":0.5}) svi = SVI(model_2, guide_2, optim, loss=Trace_ELBO()) num_samples = 50 num_iterations = 50 acc = [] for j in range(num_iterations): loss = svi.step(x_train, y_train) def predict(x): sampled_models = [guide_2(None, None) for _ in range(num_samples)] yhats = [model(x).data for model in sampled_models] mean = torch.mean(torch.stack(yhats), 0) return np.argmax(mean.numpy(), axis=1) #return mean predicted = predict(x_test) predicted = torch.tensor(predicted) total = y_test.size(0) correct = (predicted == y_test).sum().item() acc.append(100*correct/total) print(j, acc[j], loss) if acc[j] >=85: break return acc
def test_dirichlet_bernoulli(Elbo, vectorized): pyro.clear_param_store() data = torch.tensor([1.0] * 6 + [0.0] * 4) def model1(data): concentration0 = constant([10.0, 10.0]) f = pyro.sample("latent_fairness", dist.Dirichlet(concentration0))[1] for i in pyro.plate("plate", len(data)): pyro.sample("obs_{}".format(i), dist.Bernoulli(f), obs=data[i]) def model2(data): concentration0 = constant([10.0, 10.0]) f = pyro.sample("latent_fairness", dist.Dirichlet(concentration0))[1] pyro.sample( "obs", dist.Bernoulli(f).expand_by(data.shape).to_event(1), obs=data ) model = model2 if vectorized else model1 def guide(data): concentration_q = pyro.param( "concentration_q", constant([15.0, 15.0]), constraint=constraints.positive ) pyro.sample("latent_fairness", dist.Dirichlet(concentration_q)) elbo = Elbo( num_particles=7, strict_enumeration_warning=False, ignore_jit_warnings=True ) optim = Adam({"lr": 0.0005, "betas": (0.90, 0.999)}) svi = SVI(model, guide, optim, elbo) for step in range(40): svi.step(data)
def test_sample(): class Model(nn.Linear, PyroModule): def __init__(self, in_features, out_features): super().__init__(in_features, out_features) self.weight = PyroSample( lambda self: dist.Normal(0, 1) .expand([self.out_features, self.in_features]) .to_event(2)) class Guide(nn.Linear, PyroModule): def __init__(self, in_features, out_features): super().__init__(in_features, out_features) self.loc = PyroParam(torch.zeros_like(self.weight)) self.scale = PyroParam(torch.ones_like(self.weight), constraint=constraints.positive) self.weight = PyroSample( lambda self: dist.Normal(self.loc, self.scale) .to_event(2)) data = torch.randn(8) model = Model(8, 2) guide = Guide(8, 2) optim = Adam({"lr": 0.01}) svi = SVI(model, guide, optim, Trace_ELBO()) for step in range(3): svi.step(data)
def main(args): optim = Adam({"lr": 0.1}) inference = SVI(model, guide, optim, loss="ELBO") # Data is an arbitrary json-like structure with tensors at leaves. one = ng_ones(1) data = { "foo": one, "bar": [0 * one, 1 * one, 2 * one], "baz": { "noun": { "concrete": 4 * one, "abstract": 6 * one, }, "verb": 2 * one, }, } print('Step\tLoss') for step in range(args.num_epochs): if step % 100 == 0: loss = inference.step(data) print('{}\t{:0.5g}'.format(step, loss)) print('Parameters:') for name in sorted(pyro.get_param_store().get_all_param_names()): print('{} = {}'.format(name, pyro.param(name).data.cpu().numpy()))
def svi_test(): rain_prob_prior = torch.tensor(.3) my_sprinkler_prob_prior = torch.tensor(.6) neighbor_sprinkler_prob_prior = torch.tensor(.2) conditioned_lawn = pyro.condition(lawn, data={ "my_lawn": torch.tensor([1.]), "neighbor_lawn": torch.tensor([0.]) }) # guide = AutoGuide(lawn) # set up the optimizer adam_params = {"lr": 0.005, "betas": (0.90, 0.999)} optimizer = Adam(adam_params) # setup the inference algorithm svi = SVI(conditioned_lawn, lawn_guide, optimizer, loss=Trace_ELBO()) n_steps = 1000 # do gradient steps for step in range(n_steps): svi.step(rain_prob_prior, my_sprinkler_prob_prior, neighbor_sprinkler_prob_prior) if step % 100 == 0: print("step: ", step) for p in [ 'rain_prob', 'my_sprinkler_prob', 'neighbor_sprinkler_prob' ]: print(p, ": ", pyro.param(p).item())
def train(self, num_steps, lr=1e-2, restart=True, autoguide=None, use_tqdm=True): if restart: pyro.clear_param_store() if autoguide is None: autoguide = AutoMultivariateNormal else: autoguide = getattr(pyro.infer.autoguide, autoguide) self.guide = autoguide(self, init_loc_fn=init_to_mean) svi = SVI(self, guide=self.guide, optim=Adam({"lr": lr}), loss=Trace_ELBO()) loss = [] if use_tqdm: iterator = tqdm.notebook.tnrange(num_steps) else: iterator = range(num_steps) for _ in iterator: loss.append(svi.step()) return loss
def fit(self, train_loader, test_loader, num_epochs=30, test_frequency=5, stop_at_loss=100.): # setup optimizer optimizer = Adam({'lr': 1.0e-3}) # loss function elbo = JitTrace_ELBO() # setup svi svi = SVI(self.model, self.guide, optimizer, loss=elbo) for epoch in range(num_epochs): epoch_loss = torch.tensor([ svi.step(x, y) for x, y in tqdm(train_loader) ]).mean().item() / train_loader.batch_size print(f'Epoch {epoch + 1} : Loss {epoch_loss}') if epoch and epoch % test_frequency == 0: # report test diagnostics test_loss = torch.tensor([ svi.step(x, y) for x, y in tqdm(test_loader) ]).mean().item() / test_loader.batch_size print(f'*Test Loss* {test_loss}') # save to disk self.save() if int(test_loss) <= int(stop_at_loss): return self
def fit(self, dataset, batch_size, num_epochs=1, verbose=0): self.train(True) svi = SVI(self.pyro_model, self.pyro_guide, Adam({"lr": self.lr}), loss="ELBO") fit_loss = [] for epoch in range(num_epochs): epoch_loss = [] timer = time.time() # Iterate over data. for x, y in dataset.batches(batch_size): inputs = to_binary(torch.from_numpy(x).long(), (x.shape[0], self.input_dim), use_cuda=self.use_cuda) inputs = variable(self, inputs) loss = svi.step(inputs) # statistics if verbose > 1: print("Batch", len(epoch_loss), "loss:", loss, "average time:", (time.time() - timer) / float(len(epoch_loss) + 1)) epoch_loss.append(loss) if verbose > 0: print("loss =", np.mean(epoch_loss, axis=0), "time =", time.time() - timer) fit_loss.append(np.mean(epoch_loss)) return fit_loss
def main(args): pyro.set_rng_seed(0) pyro.enable_validation(__debug__) optim = Adam({"lr": 0.1}) inference = SVI(model, guide, optim, loss=Trace_ELBO()) # Data is an arbitrary json-like structure with tensors at leaves. one = torch.tensor(1.0) data = { "foo": one, "bar": [0 * one, 1 * one, 2 * one], "baz": { "noun": { "concrete": 4 * one, "abstract": 6 * one, }, "verb": 2 * one, }, } print('Step\tLoss') loss = 0.0 for step in range(args.num_epochs): loss += inference.step(data) if step and step % 10 == 0: print('{}\t{:0.5g}'.format(step, loss)) loss = 0.0 print('Parameters:') for name, value in sorted(pyro.get_param_store().items()): print('{} = {}'.format(name, value.detach().cpu().numpy()))
def init_opt(self): if self.guide is None: self.build_guide() adam_params = {"lr": 0.01, "betas": (0.90, 0.999)} optimizer = Adam(adam_params) self.svi = SVI(self.model, self.guide, optimizer, loss=Trace_ELBO())
def main(args): pyro.clear_param_store() data = build_linear_dataset(N, p) if args.cuda: # make tensors and modules CUDA data = data.cuda() softplus.cuda() regression_model.cuda() # perform inference optim = Adam({"lr": 0.05}) elbo = JitTrace_ELBO() if args.jit else Trace_ELBO() svi = SVI(model, guide, optim, loss=elbo) for j in range(args.num_epochs): if args.batch_size == N: # use the entire data set epoch_loss = svi.step(data) else: # mini batch epoch_loss = 0.0 perm = torch.randperm(N) if not args.cuda else torch.randperm( N).cuda() # shuffle data data = data[perm] # get indices of each batch all_batches = get_batch_indices(N, args.batch_size) for ix, batch_start in enumerate(all_batches[:-1]): batch_end = all_batches[ix + 1] batch_data = data[batch_start:batch_end] epoch_loss += svi.step(batch_data) if j % 100 == 0: print("epoch avg loss {}".format(epoch_loss / float(N)))
def main(args): # pyro.enable_validation(True) logging.info('Generating data') pyro.set_rng_seed(0) # We can generate synthetic data directly by calling the model. true_topic_weights, true_topic_words, data = model_original(args=args) # We'll train using SVI. logging.info('-' * 40) logging.info('Training on {} documents'.format(args.num_docs)) # wy: currently don't do enumeration. # # Elbo = JitTraceEnum_ELBO if args.jit else TraceEnum_ELBO # Elbo = TraceEnum_ELBO Elbo = TraceGraph_ELBO elbo = Elbo(max_plate_nesting=2) optim = Adam({'lr': args.learning_rate}) svi = SVI(model, guide, optim, elbo) logging.info('Step\tLoss') for step in range(args.num_steps): loss = svi.step(data, args=args, batch_size=args.batch_size) if step % 10 == 0: logging.info('{: >5d}\t{}'.format(step, loss)) loss = elbo.loss(model, guide, data, args=args) logging.info('final loss = {}'.format(loss))
def train(self, dataset, model, logger=print): m = model.model() m.to(self.device) if self.config['train']['load_model']: model.load() loss_fn = pyro.infer.Trace_ELBO(max_plate_nesting=1).differentiable_loss epochs = self.config['train']['epochs'] total_step = len(dataset) // self.config['train']['batch_size'] dataloader = DataLoader(dataset, batch_size=self.config['train']['batch_size'], shuffle=True, num_workers=4) optim = Adam({"lr": self.config['train']['learn_rate']}) svi = SVI(m.model, m.guide, optim, loss=loss_fn, num_samples=1000) n_iter = 0 for epoch in range(epochs): logger('epoch: ', epoch) for i, (X, y) in tqdm(enumerate(dataloader)): X = X.view(-1, len(self.config['dataset']['features'])) X = X.to(self.device) y = y.to(self.device) loss = svi.step(X, y) if (i+1) % 2 == 0: logger('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}' .format(epoch+1, epochs, i+1, total_step, loss/total_step)) if self.writer: self.writer.add_scalar('data/loss', loss.item(), n_iter) n_iter += 1 if epoch % self.config['train']['save_every_epoch'] == 0: model.save() model.save() if self.writer: self.writer.close()
def test_shapes(shape, stein_kernel): pyro.clear_param_store() shape1, shape2 = (5, ) + shape, shape + (6, ) mean_init1 = torch.arange( _product(shape1)).double().reshape(shape1) / 100.0 mean_init2 = torch.arange(_product(shape2)).double().reshape(shape2) def model(): pyro.sample("z1", dist.LogNormal(mean_init1, 1.0e-8).to_event(len(shape1))) pyro.sample("scalar", dist.Normal(0.0, 1.0)) pyro.sample("z2", dist.Normal(mean_init2, 1.0e-8).to_event(len(shape2))) num_particles = 7 svgd = SVGD(model, stein_kernel(), Adam({"lr": 0.0}), num_particles, 0) for step in range(2): svgd.step() particles = svgd.get_named_particles() assert particles['z1'].shape == (num_particles, ) + shape1 assert particles['z2'].shape == (num_particles, ) + shape2 for particle in range(num_particles): assert_equal(particles['z1'][particle, ...], mean_init1.exp(), prec=1.0e-6) assert_equal(particles['z2'][particle, ...], mean_init2, prec=1.0e-6)
def test_linear_regression_smoke(auto_class, Elbo): N, D = 10, 3 class RandomLinear(nn.Linear, PyroModule): def __init__(self, in_features, out_features): super().__init__(in_features, out_features) self.weight = PyroSample(dist.Normal(0., 1.).expand([out_features, in_features]).to_event(2)) self.bias = PyroSample(dist.Normal(0., 10.).expand([out_features]).to_event(1)) class LinearRegression(PyroModule): def __init__(self): super().__init__() self.linear = RandomLinear(D, 1) def forward(self, x, y=None): mean = self.linear(x).squeeze(-1) sigma = pyro.sample("sigma", dist.LogNormal(0., 1.)) with pyro.plate('plate', N): return pyro.sample('obs', dist.Normal(mean, sigma), obs=y) x, y = torch.randn(N, D), torch.randn(N) model = LinearRegression() guide = auto_class(model) infer = SVI(model, guide, Adam({'lr': 0.005}), Elbo(strict_enumeration_warning=False)) infer.step(x, y)
def train(x_data, y_data, dataset_total_length, use_cuda, n_iterations): batch_size = 512 optim = Adam({"lr": 0.005}) train_ds = TensorDataset(x_data, y_data) train_dl = DataLoader(train_ds, batch_size=batch_size, shuffle=False) kl_factor = train_dl.batch_size / len(train_dl.dataset) svi = SVI(model, guide, optim, loss=Trace_ELBO(), num_samples=1000) pyro.clear_param_store() for j in range(n_iterations): start = time.time() loss = 0 for x, y in train_dl: if use_cuda: x.cuda() y.cuda() # calculate the loss and take a gradient step loss += svi.step(dataset_total_length, x, y, kl_factor) elapsed_time = time.time() - start if j % 1 == 0: print("[iteration %04d] loss: %.4f took: %ds" % (j + 1, loss, elapsed_time)) return svi
def test_subsample_guide_2(auto_class, independent): # Simplified from Model2 in tutorial/source/forecasting_iii.ipynb def model(data): size, size = data.shape origin_plate = pyro.plate("origin", size, dim=-2) destin_plate = pyro.plate("destin", size, dim=-1) with origin_plate, destin_plate: batch = pyro.subsample(data, event_dim=0) assert batch.size(0) == batch.size(1), batch.shape pyro.sample("obs", dist.Normal(0, 1), obs=batch) def create_plates(data): size, size = data.shape origin_plate = pyro.plate("origin", size, subsample_size=5, dim=-2) if independent: destin_plate = pyro.plate("destin", size, subsample_size=5, dim=-1) else: with origin_plate as subsample: pass destin_plate = pyro.plate("destin", size, subsample=subsample, dim=-1) return origin_plate, destin_plate guide = auto_class(model, create_plates=create_plates) svi = SVI(model, guide, Adam({"lr": 0.01}), Trace_ELBO()) data = torch.randn(10, 10) for step in range(2): svi.step(data)
def VI(self, x_data, y_data, num_samples=1000, num_iterations=30000): self.guide = AutoDiagonalNormal(self.model) optim = Adam({"lr": 1e-3}) loss = Trace_ELBO() svi = SVI(self.model, self.guide, optim=optim, loss=loss) # train # print("train!") pyro.clear_param_store() # print("train2!") for j in range(num_iterations): # print("train3!") # print("x_data.size():{}, y_data.size():{}".format(x_data.size(), y_data.size())) loss = svi.step(x_data, y_data) # print("train4!") # if j % (num_iterations // 10) == 0: print("[iteration %05d] loss: %.4f" % (j + 1, loss / len(x_data))) print("sample!") # num_samplesだけ事後分布からサンプルを生成 dict = {} for i in range(num_samples): sample = self.guide() # sampling for name, value in sample.items(): if not dict.keys().__contains__(name): dict[name] = value.unsqueeze(0) else: dict[name] = torch.cat( [dict[name], value.unsqueeze(0)], dim=0) self.posterior_samples = dict
def test_em_nested_in_svi(assignment_grad): args = make_args() args.assignment_grad = assignment_grad detections = generate_data(args) pyro.clear_param_store() pyro.param('noise_scale', torch.tensor(args.init_noise_scale), constraint=constraints.positive) pyro.param('objects_loc', torch.randn(args.max_num_objects, 1)) # Learn object_loc via EM and noise_scale via SVI. optim = Adam({'lr': 0.1}) elbo = TraceEnum_ELBO(max_plate_nesting=2) newton = Newton(trust_radii={'objects_loc': 1.0}) svi = SVI(poutine.block(model, hide=['objects_loc']), poutine.block(guide, hide=['objects_loc']), optim, elbo) for svi_step in range(50): for em_step in range(2): objects_loc = pyro.param('objects_loc').detach_().requires_grad_() assert pyro.param('objects_loc').grad_fn is None loss = elbo.differentiable_loss(model, guide, detections, args) # E-step updated = newton.get_step(loss, {'objects_loc': objects_loc}) # M-step assert updated['objects_loc'].grad_fn is not None pyro.get_param_store()['objects_loc'] = updated['objects_loc'] assert pyro.param('objects_loc').grad_fn is not None loss = svi.step(detections, args) logger.debug( 'step {: >2d}, loss = {:0.6f}, noise_scale = {:0.6f}'.format( svi_step, loss, pyro.param('noise_scale').item()))
def assert_ok(model, guide, **kwargs): """ Assert that inference works without warnings or errors. """ pyro.clear_param_store() inference = SVI(model, guide, Adam({"lr": 1e-6}), "ELBO", **kwargs) inference.step()
def fit(self, optim=Adam({'lr': 1e-3}), loss=Trace_ELBO(num_particles=1), max_iter=5000, random_instance=None): """ :param optim: 优化算法 :param loss: 损失函数 :param max_iter: 最大迭代次数 :param random_instance: 随机数据生成实例 """ svi = SVI(self.model, self.guide, optim=optim, loss=loss) with trange(max_iter) as t: for i in t: t.set_description(f'迭代:{i}') svi.step(self.data) loss = svi.evaluate_loss(self.data) if isinstance(optim, PyroLRScheduler): optim.step() with torch.no_grad(): postfix_kwargs = {} if random_instance is not None: g = pyro.param('g') s = pyro.param('s') postfix_kwargs.update({ 'g': '{0}'.format((g - random_instance.g).abs().mean()), 's': '{0}'.format((s - random_instance.s).abs().mean()) }) t.set_postfix(loss=loss, **postfix_kwargs)
def train(self, epochs: int = 3000, lr: float = 0.005, window: int = 1000, log: bool = True): pyro.get_param_store().clear() adam_params = {"lr": lr, "betas": (0.95, 0.999)} optimizer = Adam(adam_params) # setup the inference algorithm svi = SVI(self.sample_model, self.sample_guide, optimizer, loss=Trace_ELBO()) n_steps = epochs # do gradient steps if log: pbar = tqdm(range(n_steps)) else: pbar = range(n_steps) elbos = [] for step in pbar: elbo = svi.step(self.x_adj, self.preds[self.mapping[0]]) elbos.append(elbo) avgs = self.ma(elbos, window) if step >= window: disp = avgs[-1] else: disp = elbo if log: pbar.set_description("Loss -> %.4f" % disp) return avgs