Automatic checkpointing

[PhilipVinc]

A long way to go...

Will need some help to make this play nice with sharding and MPI.

Must decide what to do with loggers. Right now we override the old ones. Can't really concatenate.

temporary interface.

from etils import epath
import orbax.checkpoint as ocp
import numpy as np

# Keeps a maximum of 3 checkpoints, and only saves every other step.
path = epath.Path('/tmp/nkcheck3')
checkpointer = ocp.CheckpointManager(
    path,
    options=ocp.CheckpointManagerOptions(max_to_keep=3, save_interval_steps=1),
)

import netket as nk
import optax

L = 10
g = nk.graph.Hypercube(length=L, n_dim=1, pbc=True)
hi = nk.hilbert.Spin(s=1 / 2, N=g.n_nodes)
ha = nk.operator.Ising(hilbert=hi, graph=g, h=1.0)
ma = nk.models.RBM(alpha=1, param_dtype=float)
sa = nk.sampler.MetropolisLocal(hi, n_chains=16)
op = nk.optimizer.Sgd(learning_rate=optax.linear_schedule(0.1, 0.0001, 500))
sr = nk.optimizer.SR(diag_shift=0.01)
vs = nk.vqs.MCState(sa, ma, n_samples=1008, n_discard_per_chain=10)

# Variational monte carlo driver with a variational state
gs = nk.VMC(ha, op, variational_state=vs, preconditioner=sr, checkpointer= checkpointer)

# Run the optimization for 500 iterations
gs.run(n_iter=50, out="test")

if you run it twice, the second run goes from 50 to 100 steps

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[codecov]

Codecov Report

Attention: Patch coverage is 34.78261% with 30 lines in your changes missing coverage. Please review.

Project coverage is 81.45%. Comparing base (3898574) to head (d72efb2).
Report is 22 commits behind head on master.

❗ Current head d72efb2 differs from pull request most recent head b895c1b

Please upload reports for the commit b895c1b to get more accurate results.

Files	Patch %	Lines
netket/driver/abstract_variational_driver.py	27.77%	24 Missing and 2 partials ⚠️
netket/experimental/driver/tdvp_common.py	20.00%	4 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff             @@
##           master    #1667       +/-   ##
===========================================
+ Coverage   50.50%   81.45%   +30.95%     
===========================================
  Files         312      303        -9     
  Lines       19073    18476      -597     
  Branches     2795     2726       -69     
===========================================
+ Hits         9632    15049     +5417     
+ Misses       8733     2791     -5942     
+ Partials      708      636       -72     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[inailuig]

(rebased it on master)

[PhilipVinc]

Are you using this by any chance?

I kind of abandoned the approach here because it had several issues.
I am experimenting with a different checkpointing approach, which seems to work better, but I need to gather more feedback from some testers before I finalise anything.

[inailuig]

Are you using this by any chance?

I kind of abandoned the approach here because it had several issues. I am experimenting with a different checkpointing approach, which seems to work better, but I need to gather more feedback from some testers before I finalise anything.

wanted to test it, hence the need for the rebase