Hello everyone,
I’m trying to get the MAP estimate directly during the MCMC, rather than afterward.
def _pyro_model(self, x_obs, lower_bound, upper_bound):
batch_size = x_obs.shape[0]
with pyro.plate("batch_dimensions", batch_size):
theta = pyro.sample(
"theta", dist.Uniform(lower_bound, upper_bound).to_event(1)
)
log_ratio = -self.network(x_obs, theta=theta).squeeze(-1)
with pyro.plate("batch_factors", batch_size):
pyro.factor("log_ratio_factor", log_ratio)
...
kernel = NUTS(self._pyro_model, max_tree_depth=5, jit_compile=True)
mcmc = MCMC(kernel, num_samples=num_samples, warmup_steps=warmup_steps, num_chains=1)
mcmc.run(data_batch, lower_bound, upper_bound)
samples = mcmc.get_samples()
chain_samples = samples["theta"]
...
All the methods I’ve tried to obtain the MAP after MCMC take quite a some time. I’ve tried:
- the determinist thing
- a loop on each sample to get the argmax of my network (it gives me the MLE)
- using trace from pyro.poutine
One idea was to keep the max in my function _pyro_model as I compute log_ratio. But with NUTS, this approach will consider every leap, not just the samples.
Is there really a way to do get the MAP without any additionnal time ?