Repeat MCMC sampling with more warmups without recompiling

numpyro
1

Hi,

I have defined a model which I am fitting with NUTS and in many cases I am finding that the time taken to compile the MCMC object is a significant portion of the time required to fit the model. This problem is accentuated because it frequently doesn’t converge at the first attempt and I have to increase num_warmup, which as far as I can tell means rebuilding the MCMC object before rerunning it Therefore, even with jit_model_args=True it still requires recompiling the model before it can be run

My question is, does numpyro provide a way to increase the number of warmup steps of an MCMC object without requiring the model to be recompiled before it can be run? I cannot see anything in the documentation describing this, but I would be surprised if I was the first person to wish to do this so I was wondering if there is anything I have missed?

Thanks in advance

1 Like
2

I think it would be tricky to allow dynamic warmup steps because currently, the adaptation schedule is tie to a specific number of steps. It is possible to support dynamic num_warmup though so please make a FR in github for this.

3

Thanks for your response @fehiepsi.

Just to clarify, would making num_warmup “dynamic” mean that we could repeatedly calll MCMC.run()
with different numbers of warmup steps without having to recompile each time?

(Presumably this would need jit_model_args=True as well, but that is fine for my use case.)

4

Yes, that’s right.

5

Thanks, I’ve opened an issue regarding this on github: Request for dynamic `num_warmup` argument to MCMC object · Issue #1211 · pyro-ppl/numpyro · GitHub

1 Like