def loss(self, model, guide, *args, **kwargs): """ :returns: returns an estimate of the ELBO :rtype: float Evaluates the ELBO with an estimator that uses num_particles many samples/particles. """ elbo = 0.0 for weight, model_trace, guide_trace, log_r in self._get_traces( model, guide, *args, **kwargs): elbo_particle = weight * 0 log_pdf = "batch_log_pdf" if ( self.enum_discrete and weight.size(0) > 1) else "log_pdf" for name in model_trace.nodes.keys(): if model_trace.nodes[name]["type"] == "sample": if model_trace.nodes[name]["is_observed"]: elbo_particle += model_trace.nodes[name][log_pdf] else: elbo_particle += model_trace.nodes[name][log_pdf] elbo_particle -= guide_trace.nodes[name][log_pdf] # drop terms of weight zero to avoid nans if isinstance(weight, numbers.Number): if weight == 0.0: elbo_particle = torch_zeros_like(elbo_particle) else: elbo_particle[weight == 0] = 0.0 elbo += torch_data_sum(weight * elbo_particle) loss = -elbo return loss
def sample(self): """ Returns a sample which has the same shape as `ps`, except that the last dimension will have the same size as the number of events. :return: sample from the OneHotCategorical distribution :rtype: torch.Tensor """ sample = torch_multinomial(self.ps.data, 1, replacement=True).expand(*self.shape()) sample_one_hot = torch_zeros_like(self.ps.data).scatter_(-1, sample, 1) return Variable(sample_one_hot)
def test_batch_log_pdf(lognormal): dist_params = lognormal.get_dist_params(0) mu_lognorm = dist_params['mu'] sigma_lognorm = dist_params['sigma'] mu_z = torch_zeros_like(mu_lognorm) sigma_z = torch_ones_like(sigma_lognorm) trans_dist = get_transformed_dist(dist.normal, sigma_lognorm, mu_lognorm) test_data = lognormal.get_test_data(0) log_px_torch = trans_dist.batch_log_pdf(test_data, mu_z, sigma_z).data.cpu().numpy() log_px_np = sp.lognorm.logpdf( test_data.data.cpu().numpy(), sigma_lognorm.data.cpu().numpy(), scale=np.exp(mu_lognorm.data.cpu().numpy())).sum(-1, keepdims=True) assert_equal(log_px_torch, log_px_np, prec=1e-4)
def batch_log_pdf(self, x): """ Evaluates log probability densities for one or a batch of samples and parameters. The last dimension for `ps` encodes the event probabilities, and the remaining dimensions are considered batch dimensions. `ps` and `vs` are first broadcasted to the size of the data `x`. The data tensor is used to to create a mask over `vs` where a 1 in the mask indicates that the corresponding value in `vs` was selected. Since, `ps` and `vs` have the same size, this mask when applied over `ps` gives the probabilities of the selected events. The method returns the logarithm of these probabilities. :return: tensor with log probabilities for each of the batches. :rtype: torch.autograd.Variable """ logits = self.logits vs = self.vs x = self._process_data(x) batch_pdf_shape = self.batch_shape(x) + (1, ) # probability tensor mask when data is numpy if isinstance(x, np.ndarray): batch_vs_size = x.shape[:-1] + (vs.shape[-1], ) vs = np.broadcast_to(vs, batch_vs_size) boolean_mask = torch.from_numpy((vs == x).astype(int)) # probability tensor mask when data is pytorch tensor else: x = x.cuda() if logits.is_cuda else x.cpu() batch_ps_shape = self.batch_shape(x) + self.event_shape() logits = logits.expand(batch_ps_shape) if vs is not None: vs = vs.expand(batch_ps_shape) boolean_mask = (vs == x) elif self.one_hot: boolean_mask = x else: boolean_mask = torch_zeros_like(logits.data).scatter_( -1, x.data.long(), 1) boolean_mask = boolean_mask.cuda( ) if logits.is_cuda else boolean_mask.cpu() if not isinstance(boolean_mask, Variable): boolean_mask = Variable(boolean_mask) # apply log function to masked probability tensor batch_log_pdf = logits.masked_select( boolean_mask.byte()).contiguous().view(batch_pdf_shape) if self.log_pdf_mask is not None: batch_log_pdf = batch_log_pdf * self.log_pdf_mask return batch_log_pdf
def test_mean_and_var(lognormal): dist_params = lognormal.get_dist_params(0) mu_lognorm = dist_params['mu'] sigma_lognorm = dist_params['sigma'] mu_z = torch_zeros_like(mu_lognorm) sigma_z = torch_ones_like(sigma_lognorm) trans_dist = get_transformed_dist(dist.normal, sigma_lognorm, mu_lognorm) torch_samples = trans_dist.sample(mu_z, sigma_z, batch_size=lognormal.get_num_samples(0)) torch_mean = torch.mean(torch_samples, 0) torch_std = torch.std(torch_samples, 0) analytic_mean = lognormal.pyro_dist.analytic_mean(**dist_params) analytic_std = lognormal.pyro_dist.analytic_var(**dist_params) ** 0.5 precision = analytic_mean.max().data[0] * 0.05 assert_equal(torch_mean, analytic_mean, prec=precision) assert_equal(torch_std, analytic_std, prec=precision)
def sample(self): """ Returns a sample which has the same shape as `ps` (or `vs`). The type of the sample is `numpy.ndarray` if `vs` is a list or a numpy array, else a tensor is returned. :return: sample from the Categorical distribution :rtype: numpy.ndarray or torch.LongTensor """ sample = torch_multinomial(self.ps.data, 1, replacement=True).expand(*self.shape()) sample_one_hot = torch_zeros_like(self.ps.data).scatter_(-1, sample, 1) if self.vs is not None: if isinstance(self.vs, np.ndarray): sample_bool_index = sample_one_hot.cpu().numpy().astype(bool) return self.vs[sample_bool_index].reshape(*self.shape()) else: return self.vs.masked_select(sample_one_hot.byte()) return Variable(sample)
def loss_and_grads(self, model, guide, *args, **kwargs): """ :returns: returns an estimate of the ELBO :rtype: float Computes the ELBO as well as the surrogate ELBO that is used to form the gradient estimator. Performs backward on the latter. Num_particle many samples are used to form the estimators. """ elbo = 0.0 # grab a trace from the generator for weight, model_trace, guide_trace, log_r in self._get_traces( model, guide, *args, **kwargs): elbo_particle = weight * 0 surrogate_elbo_particle = weight * 0 # compute elbo and surrogate elbo log_pdf = "batch_log_pdf" if ( self.enum_discrete and weight.size(0) > 1) else "log_pdf" for name in model_trace.nodes.keys(): if model_trace.nodes[name]["type"] == "sample": if model_trace.nodes[name]["is_observed"]: elbo_particle += model_trace.nodes[name][log_pdf] surrogate_elbo_particle += model_trace.nodes[name][ log_pdf] else: lp_lq = model_trace.nodes[name][ log_pdf] - guide_trace.nodes[name][log_pdf] elbo_particle += lp_lq if guide_trace.nodes[name]["fn"].reparameterized: surrogate_elbo_particle += lp_lq else: # XXX should the user be able to control inclusion of the -logq term below? surrogate_elbo_particle += model_trace.nodes[name][log_pdf] + \ log_r.detach() * guide_trace.nodes[name][log_pdf] # drop terms of weight zero to avoid nans if isinstance(weight, numbers.Number): if weight == 0.0: elbo_particle = torch_zeros_like(elbo_particle) surrogate_elbo_particle = torch_zeros_like( surrogate_elbo_particle) else: weight_eq_zero = (weight == 0) elbo_particle[weight_eq_zero] = 0.0 surrogate_elbo_particle[weight_eq_zero] = 0.0 elbo += torch_data_sum(weight * elbo_particle) surrogate_elbo_particle = torch_sum(weight * surrogate_elbo_particle) # collect parameters to train from model and guide trainable_params = set(site["value"] for trace in (model_trace, guide_trace) for site in trace.nodes.values() if site["type"] == "param") if trainable_params: surrogate_loss_particle = -surrogate_elbo_particle torch_backward(surrogate_loss_particle) pyro.get_param_store().mark_params_active(trainable_params) loss = -elbo return loss