Resuming sampling from a previous trace

Hi. Is there a way to resume sampling from a previous trace / inference data (or other state from a previous sampling run)? I’d like to be able to draw more samples, without having to draw any tuning samples.

I’m using the NUTS sampler and I’ve tried the approach @junpenglao outlines in this discussion, but the methods he outlines fail for me:

  • The new model approach fails because models no longer have a bijection attribute
  • The same model approach draws identical samples that do not change

Is there a way to resuming sampling from a previous trace / inference data? Ideally, I’m looking for a solution that I could use for any sampling methods, not just NUTS. If not, is this something that pymc is planning to support? It seems like a fairly common use case would be to draw some samples, analyze the statistics / convergence statistics and then draw more samples

Thanks,
Laurence

Welcome!

You might want to check out related discussion here.

Hi @cluhmann. Thanks for the link. I’ve tried the solutions outlined in that discussion, but they didn’t work:

  • The iter_sample function no longer exists
  • If I use mcbackend to draw samples, then use it to draw samples without tuning, every sample in the second set of draws is identical

Interesting. Maybe @michaelosthege has some suggestions?

Essentially this is a matter of the PyMC samplers being stateful and there currently not being a standardized API to save/restore that state.
This includes things like instance attributes, but also random states (see Refactor step methods to use their own random stream · Issue #5797 · pymc-devs/pymc · GitHub).

I’m currently refactoring the trace backend to be natively McBackend-compatible, aiming to eventually delete pm.backends.BaseTrace, pm.backends.ndarray.NDArray in favor of defaulting to mcbackend.NumPyBackend which supports sparse sampler stats (Support sparse sample stats · Issue #6194 · pymc-devs/pymc · GitHub).

With sparse sampler stats we can start saving the sampler state in sampler stats, for example by storing the mass matrix information in NUTS as a sampler stat every time it changed during tuning, or by emitting the current random state as a sampler stat every 100 iterations or so.
Having such “keyframes” in the trace could then be the starting point for properly restoring stateful samplers and resuming an MCMC.

With Refactoring towards `IBaseTrace` interfaces by michaelosthege · Pull Request #6475 · pymc-devs/pymc · GitHub I’m actually getting pretty close to optional McBackend support already. Any help with refactoring the step method interface (e.g. inclucing shape information in .stats_dtypes, or taking care of Refactor step methods to use their own random stream · Issue #5797 · pymc-devs/pymc · GitHub) would be greatly appreciated!

1 Like