def forward(self, t): # t, yt: (n,) n = len(t) dt = t[1] - t[0] nzero = torch.zeros(n) # sample the prior a = pyro.sample("a", Uniform(*self.a_bounds)) log_lscale = pyro.sample("log_lscale", Uniform(*self.log_lscale_bounds)) lscale = torch.exp(log_lscale) log_sigma = pyro.sample("log_sigma", Uniform(*self.logs_bounds)) tdist = t.unsqueeze(-1) - t # (n,n) b_sigma = pyro.sample("b_sigma", Uniform(*self.b_bounds)) b_cov = b_sigma * b_sigma * torch.exp( -tdist * tdist / (2 * lscale * lscale)) + torch.eye(n) * 1e-5 b = pyro.sample("b", MultivariateNormal(nzero, b_cov)) # calculate the rate int_bdt = torch.cumsum(b, dim=0) * dt mu = a + int_bdt # (n,) # simulate the observation logysim = pyro.sample("logyt", Normal(mu, torch.exp(log_sigma))) return logysim
def test_add(): X = Uniform(0, 1).rv # (0, 1) X = X + 1 # (1, 2) X = 1 + X # (2, 3) X += 1 # (3, 4) x = X.dist.sample([N_SAMPLES]) assert ((3 <= x) & (x <= 4)).all().item()
def test_subtract(): X = Uniform(0, 1).rv # (0, 1) X = 1 - X # (0, 1) X = X - 1 # (-1, 0) X -= 1 # (-2, -1) x = X.dist.sample([N_SAMPLES]) assert ((-2 <= x) & (x <= -1)).all().item()
def _model(x_data, y_data): # weight and bias priors w_prior = Normal(torch.zeros(1, 1), torch.ones(1, 1)).to_event(1) b_prior = Normal(10 * torch.ones(1, 1), 10 * torch.ones(1, 1)).to_event(1) priors = {'linear.weight': w_prior, 'linear.bias': b_prior} scale = pyro.sample('sigma', Uniform(0, 2000)) lifted_module = pyro.random_module("module", regression_model, priors) # sample a nn (which also samples w and b) lifted_reg_model = lifted_module() with pyro.plate("map", len(x_data)): # run the nn forward on data prediction_mean = lifted_reg_model(x_data).squeeze( -1) # shape: (256,) # condition on the observed data res = pyro.sample("obs", Normal(prediction_mean, scale), obs=y_data) # shape (256, 1) return prediction_mean
def _model(x_data, y_data): fc1w_prior = Normal(torch.zeros(2, 3), torch.ones(2, 3)).to_event() fc1b_prior = Normal(0.25 * torch.ones(2), 0.25 * torch.ones(2)).to_event() outw_prior = Normal(torch.zeros(1, 2), torch.ones(1, 2)).to_event() outb_prior = Normal(0.25 * torch.ones(1), 0.25 * torch.ones(1)).to_event() priors = { 'fc1.weight': fc1w_prior, 'fc1.bias': fc1b_prior, 'out.weight': outw_prior, 'out.bias': outb_prior } scale = pyro.sample('sigma', Uniform(0, 20)) # lift module parameters to random variables sampled from the priors lifted_module = pyro.random_module("module", nn_model, priors) # sample a regressor (which also samples w and b) lifted_reg_model = lifted_module() with pyro.plate("map", len(x_data)): # run the nn forward on data prediction_mean = lifted_reg_model(x_data).squeeze(-1) # condition on the observed data pyro.sample("obs", Normal(prediction_mean, scale), obs=y_data) return prediction_mean
def pyromodel(x, y): priors = {} for name, par in model.named_parameters(): priors[name] = dist.Normal(torch.zeros(*par.shape), 50 * torch.ones(*par.shape)).independent( par.dim()) #print("batch shape:", priors[name].batch_shape) #print("event shape:", priors[name].event_shape) #print("event dim:", priors[name].event_dim) bayesian_model = pyro.random_module('bayesian_model', model, priors) sampled_model = bayesian_model() sigma = pyro.sample('sigma', Uniform(0, 50)) with pyro.iarange("map", len(x)): prediction_mean = sampled_model(x) logging.debug(f"prediction_mean: {prediction_mean.shape}") if y is not None: logging.debug(f"y_data: {y.shape}") d_dist = Normal(prediction_mean, sigma).to_event(1) if y is not None: logging.debug(f"y_data: {y.shape}") logging.debug(f"batch shape: {d_dist.batch_shape}") logging.debug(f"event shape: {d_dist.event_shape}") logging.debug(f"event dim: {d_dist.event_dim}") pyro.sample("obs", d_dist, obs=y) return prediction_mean
def guide_t0(data): # T-1 alpha params for beta sampling kappa = pyro.param('kappa', lambda: Uniform(0, 2).sample([T - 1]), constraint=constraints.positive) # concentration params for q_theta #[T,C] tau = pyro.param('tau', lambda: MultivariateNormal(0.5 * torch.ones(C), 0.25 * torch.eye(C)).sample([T]), constraint=constraints.unit_interval) # N params for categorical dist; topic weights; symmetric prior phi = pyro.param('phi', lambda: Dirichlet(1 / T * torch.ones(T)).sample([N]), constraint=constraints.simplex) with pyro.plate("beta_plate", T - 1): q_beta = 0 q_beta += pyro.sample("beta", Beta(torch.ones(T - 1), kappa)) # q_beta *= 1 # sample probs for multinomial distributions with pyro.plate("theta_plate", T): # outputs multinomial probabilities for each topic q_theta = 0 q_theta += pyro.sample("theta", Dirichlet(tau)) # q_theta *= 1 with pyro.plate("data", N): z = 0 z += pyro.sample("z", Categorical(phi))
def _traces(self, *args, **kwargs): """ make trace posterior distribution """ # initialize traces with a draw from the prior old_model_trace = poutine.trace(self.model)(*args, **kwargs) traces = [] t = 0 i = 0 while t < self.burn + self.lag * self.samples: i += 1 # q(z' | z) new_guide_trace = poutine.block(poutine.trace(self.model))( old_model_trace, *args, **kwargs) # p(x, z') new_model_trace = poutine.trace( poutine.replay(self.model, new_guide_trace))(*args, **kwargs) # q(z | z') old_guide_trace = poutine.block( poutine.trace(poutine.replay(self.model, old_model_trace)))( new_model_trace, *args, **kwargs) # p(x, z') q(z' | z) / p(x, z) q(z | z') logr = new_model_trace.log_pdf() + new_guide_trace.log_pdf() - \ old_model_trace.log_pdf() - old_guide_trace.log_pdf() rnd = pyro.sample("mh_step_{}".format(i), Uniform(torch.zeros(1), torch.ones(1))) if torch.log(rnd).data[0] < logr.data[0]: # accept t += 1 old_model_trace = new_model_trace if t <= self.burn or (t > self.burn and t % self.lag == 0): yield (new_model_trace, new_model_trace.log_pdf())
def model(x_data, y_data): fc1w_prior = Normal(loc=torch.zeros_like(net.fc1.weight), scale=torch.ones_like(net.fc1.weight)) fc1b_prior = Normal(loc=torch.zeros_like(net.fc1.bias), scale=torch.ones_like(net.fc1.bias)) outw_prior = Normal(loc=torch.zeros_like(net.out.weight), scale=torch.ones_like(net.out.weight)) outb_prior = Normal(loc=torch.zeros_like(net.out.bias), scale=torch.ones_like(net.out.bias)) priors = {'fc1.weight': fc1w_prior, 'fc1.bias': fc1b_prior, 'out.weight': outw_prior, 'out.bias': outb_prior} f_prior = Normal(0., 1.) priors = {'linear.weight': w_prior, 'linear.bias': b_prior, 'factor': f_prior} scale = pyro.sample("sigma", Uniform(0., 10.)) # lift module parameters to random variables sampled from the priors lifted_module = pyro.random_module("module", net, priors) # sample a regressor (which also samples w and b) lifted_reg_model = lifted_module() """ lhat = log_softmax(lifted_reg_model(x_data)) pyro.sample("obs", Categorical(logits=lhat), obs=y_data) """ # run the regressor forward conditioned on inputs prediction_mean = lifted_reg_model(x_data).squeeze(-1) pyro.sample("obs", Normal(prediction_mean, scale), obs=y_data) return prediction_mean
def model(x_data, y_data): # weight and bias priors fc1w_prior = Normal(torch.zeros(1, 2), torch.ones(1, 2)).to_event(1) fc1b_prior = Normal(torch.tensor([[8.]]), torch.tensor([[1000.]])).to_event(1) outw_prior = Normal(loc=torch.zeros_like(net.out.weight), scale=torch.ones_like(net.out.weight)) outb_prior = Normal(loc=torch.zeros_like(net.out.bias), scale=torch.ones_like(net.out.bias)) #outw_prior = Normal(loc=outw_mu_param, scale=outw_sigma_param).independent(1) #outw_prior = Normal(loc=outw_mu_param, scale=outw_sigma_param).independent(1) #f_prior = Normal(0., 1.) #priors = {'linear.weight': w_prior, 'linear.bias': b_prior, 'factor': f_prior} priors = { 'fc1.weight': fc1w_prior, 'fc1.bias': fc1b_prior, 'out.weight': outw_prior, 'out.bias': outb_prior } scale = pyro.sample("sigma", Uniform(0., 10.)) # lift module parameters to random variables sampled from the priors lifted_module = pyro.random_module("module", net, priors) # sample a nn (which also samples w and b) lifted_reg_model = lifted_module() with pyro.plate("map", len(x_data)): # run the nn forward on data prediction_mean = lifted_reg_model(x_data).squeeze(-1) # condition on the observed data pyro.sample("obs", Normal(prediction_mean, scale), obs=y_data) return prediction_mean
def model(x_data, y_data, regression_model): p = x_data.shape[1] # weight and bias priors # w_prior = Normal(torch.zeros(1, 2), torch.ones(1, 2)).to_event(1) # b_prior = Normal(torch.tensor([[8.]]), torch.tensor([[1000.]])).to_event(1) w_prior = Normal(torch.zeros(1, p), torch.ones(1, p)).to_event(1) b_prior = Normal(torch.tensor([[1.]]), torch.tensor([[10.]])).to_event(1) f_prior = Normal(0., 1.) priors = { 'linear.weight': w_prior, 'linear.bias': b_prior, 'factor': f_prior } scale = pyro.sample("sigma", Uniform(0., 10.)) # lift module parameters to random variables sampled from the priors lifted_module = pyro.random_module("module", regression_model, priors) # sample a nn (which also samples w and b) lifted_reg_model = lifted_module() with pyro.plate("map", len(x_data)): # run the nn forward on data prediction_mean = lifted_reg_model(x_data).squeeze(-1) # condition on the observed data pyro.sample("obs", Normal(prediction_mean, scale), obs=y_data) return prediction_mean
def model(self, features, target): def normal_prior(x): return Normal(torch.zeros_like(x), torch.ones_like(x)).to_event(x.dim()) self.priors = {} for i in range(len(self.net.hidden_sizes)): self.priors['h' + str(i) + '.weight'] = normal_prior( getattr(self.net, 'h' + str(i)).weight) self.priors['h' + str(i) + '.bias'] = normal_prior( getattr(self.net, 'h' + str(i)).bias) self.priors['out' + '.weight'] = normal_prior(self.net.out.weight) self.priors['out' + '.bias'] = normal_prior(self.net.out.bias) # lift module parameters to random variables sampled from the priors lifted_module = pyro.random_module("module", self.net, self.priors) # sample a regressor (which also samples w and b) model_sample = lifted_module() out_sigma = pyro.sample("sigma", Uniform(0., 10.)) # precision = pyro.sample("precision", Uniform(0., 10.)) # out_sigma = 1 / precision with pyro.plate("data", len(target)): target_mean = model_sample(features).squeeze(-1) # target is not one-hot encoded pyro.sample("obs", Normal(target_mean, out_sigma), obs=target) return target_mean
def model(is_cont_africa, ruggedness, log_gdp): a = pyro.sample("a", Normal(8., 1000.)) b_a = pyro.sample("bA", Normal(0., 1.)) b_r = pyro.sample("bR", Normal(0., 1.)) b_ar = pyro.sample("bAR", Normal(0., 1.)) sigma = pyro.sample("sigma", Uniform(0., 10.)) mean = a + b_a * is_cont_africa + b_r * ruggedness + b_ar * is_cont_africa * ruggedness with pyro.plate("data", 170): pyro.sample("obs", Normal(mean, sigma), obs=log_gdp)
def sample(self): # TODO: write a better implementation f = lambda a: Uniform(a[0], a[1]).sample() lis = list(self.condition_matrix.size()) sample = torch.zeros((lis[0], lis[1])) for i in range(lis[0]): for j in range(lis[1]): sample[i, j] += f(self.condition_matrix[i, j]) return sample
def test_tensor_ops(): pi = 3.141592654 X = Uniform(0, 1).expand([5, 5]).rv a = tt([[1, 2, 3, 4, 5]]) b = a.T X = abs(pi*(-X + a - 3*b)) x = X.dist.sample() assert x.shape == (5, 5) assert (x >= 0).all().item()
def test_clippedadam_lrd(lrd): x1 = torch.tensor(0.0, requires_grad=True) orig_lr = 1.0 opt_ca = optim.clipped_adam.ClippedAdam(params=[x1], lr=orig_lr, lrd=lrd) for step in range(3): g = Uniform(-5.0, 5.0).sample() x1.backward(g) opt_ca.step() assert opt_ca.param_groups[0]["lr"] == orig_lr * lrd**(step + 1)
def survives(self, t, λ, μ, ρ): t_end = t - Exponential(μ).sample() if t_end <= 0: if Bernoulli(ρ).sample(): return True t_end = 0 for i in range(int(Poisson(λ * (t - t_end)).sample())): τ = Uniform(t_end, t).sample() if self.survives(τ, λ, μ, ρ): return True return False
def fully_pooled(at_bats): """ Number of hits in $K$ at bats for each player has a Binomial distribution with a common probability of success, $\phi$. :param (torch.Tensor) at_bats: Number of at bats for each player. :return: Number of hits predicted by the model. """ phi_prior = Uniform(at_bats.new_tensor(0), at_bats.new_tensor(1)) phi = pyro.sample("phi", phi_prior) return pyro.sample("obs", Binomial(at_bats, phi))
def not_pooled(at_bats): """ Number of hits in $K$ at bats for each player has a Binomial distribution with independent probability of success, $\phi_i$. :param (torch.Tensor) at_bats: Number of at bats for each player. :return: Number of hits predicted by the model. """ num_players = at_bats.shape[0] phi_prior = Uniform(at_bats.new_tensor(0), at_bats.new_tensor(1)).expand_by([num_players]).independent(1) phi = pyro.sample("phi", phi_prior) return pyro.sample("obs", Binomial(at_bats, phi))
def test_chaining(): X = ( Uniform(0, 1).rv # (0, 1) .add(1) # (1, 2) .pow(2) # (1, 4) .mul(2) # (2, 8) .sub(5) # (-3, 3) .tanh() # (-1, 1); more like (-0.995, +0.995) .exp() # (1/e, e) ) x = X.dist.sample([N_SAMPLES]) assert ((1/math.e <= x) & (x <= math.e)).all().item()
def test_clippedadam_clip(clip_norm): x1 = torch.tensor(0., requires_grad=True) x2 = torch.tensor(0., requires_grad=True) opt_ca = optim.clipped_adam.ClippedAdam(params=[x1], lr=1., lrd=1., clip_norm=clip_norm) opt_a = torch.optim.Adam(params=[x2], lr=1.) for step in range(3): opt_ca.zero_grad() opt_a.zero_grad() x1.backward(Uniform(clip_norm, clip_norm + 3.).sample()) x2.backward(torch.tensor(clip_norm)) opt_ca.step() opt_a.step() assert_equal(x1, x2)
def model(observed_data): mu_prior = Uniform(0.0, 1.0) mu = pyro.sample("mu", mu_prior) def observe_T(T_obs, obs_name): T_simulated = simulate(mu) T_obs_dist = Normal(T_simulated, torch.tensor(time_measurement_sigma)) pyro.sample(obs_name, T_obs_dist, obs=T_obs) for i, T_obs in enumerate(observed_data): observe_T(T_obs, "obs_%d" % i) return mu
def pgm_model(self): sex_dist = Bernoulli(logits=self.sex_logits).to_event(1) # pseudo call to register with pyro _ = self.sex_logits sex = pyro.sample('sex', sex_dist, infer=dict(baseline={'use_decaying_avg_baseline': True})) slice_number_dist = Uniform(self.slice_number_min, self.slice_number_max).to_event(1) slice_number = pyro.sample('slice_number', slice_number_dist) age_base_dist = Normal(self.age_base_loc, self.age_base_scale).to_event(1) age_dist = TransformedDistribution(age_base_dist, self.age_flow_transforms) _ = self.age_flow_components age = pyro.sample('age', age_dist) age_ = self.age_flow_constraint_transforms.inv(age) duration_context = torch.cat([sex, age_], 1) duration_base_dist = Normal(self.duration_base_loc, self.duration_base_scale).to_event(1) duration_dist = ConditionalTransformedDistribution(duration_base_dist, self.duration_flow_transforms).condition(duration_context) # noqa: E501 duration = pyro.sample('duration', duration_dist) _ = self.duration_flow_components duration_ = self.duration_flow_constraint_transforms.inv(duration) edss_context = torch.cat([sex, duration_], 1) edss_base_dist = Normal(self.edss_base_loc, self.edss_base_scale).to_event(1) edss_dist = ConditionalTransformedDistribution(edss_base_dist, self.edss_flow_transforms).condition(edss_context) # noqa: E501 edss = pyro.sample('edss', edss_dist) _ = self.edss_flow_components edss_ = self.edss_flow_constraint_transforms.inv(edss) brain_context = torch.cat([sex, age_], 1) brain_volume_base_dist = Normal(self.brain_volume_base_loc, self.brain_volume_base_scale).to_event(1) brain_volume_dist = ConditionalTransformedDistribution(brain_volume_base_dist, self.brain_volume_flow_transforms).condition(brain_context) _ = self.brain_volume_flow_components brain_volume = pyro.sample('brain_volume', brain_volume_dist) brain_volume_ = self.brain_volume_flow_constraint_transforms.inv(brain_volume) ventricle_context = torch.cat([age_, brain_volume_, duration_], 1) ventricle_volume_base_dist = Normal(self.ventricle_volume_base_loc, self.ventricle_volume_base_scale).to_event(1) ventricle_volume_dist = ConditionalTransformedDistribution(ventricle_volume_base_dist, self.ventricle_volume_flow_transforms).condition(ventricle_context) # noqa: E501 ventricle_volume = pyro.sample('ventricle_volume', ventricle_volume_dist) _ = self.ventricle_volume_flow_components ventricle_volume_ = self.ventricle_volume_flow_constraint_transforms.inv(ventricle_volume) lesion_context = torch.cat([brain_volume_, ventricle_volume_, duration_, edss_], 1) lesion_volume_base_dist = Normal(self.lesion_volume_base_loc, self.lesion_volume_base_scale).to_event(1) lesion_volume_dist = ConditionalTransformedDistribution(lesion_volume_base_dist, self.lesion_volume_flow_transforms).condition(lesion_context) lesion_volume = pyro.sample('lesion_volume', lesion_volume_dist) _ = self.lesion_volume_flow_components return dict(age=age, sex=sex, ventricle_volume=ventricle_volume, brain_volume=brain_volume, lesion_volume=lesion_volume, duration=duration, edss=edss, slice_number=slice_number)
def model(observed_data): mu_prior = Uniform(Variable(torch.zeros(1)), Variable(torch.ones(1))) mu = pyro.sample("mu", mu_prior) def observe_T(T_obs, obs_name): T_simulated = simulate(mu) T_obs_dist = Normal(T_simulated, Variable(torch.Tensor([time_measurement_sigma]))) pyro.observe(obs_name, T_obs_dist, T_obs) for i, T_obs in enumerate(observed_data): observe_T(T_obs, "obs_%d" % i) return mu
def not_pooled(at_bats, hits): r""" Number of hits in $K$ at bats for each player has a Binomial distribution with independent probability of success, $\phi_i$. :param (torch.Tensor) at_bats: Number of at bats for each player. :param (torch.Tensor) hits: Number of hits for the given at bats. :return: Number of hits predicted by the model. """ num_players = at_bats.shape[0] with pyro.plate("num_players", num_players): phi_prior = Uniform(scalar_like(at_bats, 0), scalar_like(at_bats, 1)) phi = pyro.sample("phi", phi_prior) return pyro.sample("obs", Binomial(at_bats, phi), obs=hits)
def test_clip_norm(pyro_optim, clip, value): x1 = torch.tensor(0.0, requires_grad=True) x2 = torch.tensor(0.0, requires_grad=True) opt_c = pyro_optim({"lr": 1.0}, {clip: value}) opt = pyro_optim({"lr": 1.0}) for step in range(3): x1.backward(Uniform(value, value + 3.0).sample()) x2.backward(torch.tensor(value)) opt_c([x1]) opt([x2]) assert_equal(x1.grad, torch.tensor(value)) assert_equal(x2.grad, torch.tensor(value)) assert_equal(x1, x2) opt_c.optim_objs[x1].zero_grad() opt.optim_objs[x2].zero_grad()
def get_new_kernel(self, thetas: Tensor) -> Distribution: """Return new kernel distribution for a given set of paramters.""" if self.kernel == "gaussian": assert self.kernel_variance.ndim == 2 return MultivariateNormal(loc=thetas, covariance_matrix=self.kernel_variance) elif self.kernel == "uniform": low = thetas - self.kernel_variance high = thetas + self.kernel_variance # Move batch shape to event shape to get Uniform that is multivariate in # parameter dimension. return Uniform(low=low, high=high).to_event(1) else: raise ValueError(f"Kernel, '{self.kernel}' not supported.")
def _skewness(event_shape): skewness = torch.zeros(event_shape.numel()) done = False while not done: for i in range(event_shape.numel()): max_ = 1.0 - skewness.abs().sum(-1) if torch.any(max_ < 1e-15): break skewness[i] = Uniform(-max_, max_).sample() done = not torch.any(max_ < 1e-15) if event_shape == tuple(): skewness = skewness.reshape(event_shape) else: skewness = skewness.view(event_shape) return skewness
def test_clippedadam_pass(clip_norm): x1 = torch.tensor(0.0, requires_grad=True) x2 = torch.tensor(0.0, requires_grad=True) opt_ca = optim.clipped_adam.ClippedAdam(params=[x1], lr=1.0, lrd=1.0, clip_norm=clip_norm) opt_a = torch.optim.Adam(params=[x2], lr=1.0) for step in range(3): g = Uniform(-clip_norm, clip_norm).sample() opt_ca.zero_grad() opt_a.zero_grad() x1.backward(g) x2.backward(g) opt_ca.step() opt_a.step() assert_equal(x1, x2)
def pyromodel(x, y): priors = {} for name, par in model.named_parameters(): priors[name] = dist.Normal(torch.zeros(*par.shape), 10 * torch.ones(*par.shape)).independent(par.dim()) #print("batch shape:", priors[name].batch_shape) #print("event shape:", priors[name].event_shape) #print("event dim:", priors[name].event_dim) bayesian_model = pyro.random_module('bayesian_model', model, priors) sampled_model = bayesian_model() sigma = pyro.sample('sigma', Uniform(0, 10)) with pyro.iarange("map", len(x)): prediction_mean = sampled_model(x).squeeze(-1) pyro.sample("obs", dist.Normal(prediction_mean, 0.05 * sigma), obs=y)