Serialised MCState (sometimes) does not produce correct results?

[waleed-sh]

Hi!

In my optimisation process, I have created a callback which checks the current loss and if it under a certain threshold, it exports the serialised MCState, so something like

    def __call__(self, step, log_data, driver):
        # get the loss from the log_data and index the most recent value
        loss = np.real(getattr(log_data[driver._loss_name], self.monitor))

        if loss <= self.target:
            # export the state here

        return True

where I have manager to pass the variational state somehow to the callback and then use the the serialisation from here.

Later, to load the state, I create the same network, same sampler, etc. and then load the state by using this function. However, when I compute the expectation value of the same operator I was minimising, I do not get the same result that I had when I saved it. This is not always the case, but it happens so often that it might as well be and I really start to question my sanity now...

Is there any reason for this to be the case? Am I doing something incorrectly? I am not sure if it is a serialisation issue or is it some issue in the order of doing things (e.g. is the callback exporting the state after the parameter update) or is it some conceptual thing (e.g. I would never get the same cost because the minimisation isnt completed yet). My goal is to basically export a "frozen state" of the current entire process that enables me to produce the same cost and look at the same variational state at the time that cost was produced. Any help would be appreciated!

P.S.
I noticed that in the serialisation process, the sometimes the previous sampler state is saved if the samples is not None, and sometimes we use the current sampler state, wouldn't we want to save the most recent sampler state that produced that loss?

[waleed-sh]

I am closing this discussion as I have discovered the error which was on my part, it was an issue of the state of the variational state being passed to the callback...

[PhilipVinc]

Serialisation of variational states is discussed in https://netket.readthedocs.io/en/v3.6/docs/varstate.html#saving-and-loading-a-variational-state

You callback could be something as simple as:

    def __call__(self, step, log_data, driver):
        # get the loss from the log_data and index the most recent value
        loss = np.real(getattr(log_data[driver._loss_name], self.monitor))

        if loss <= self.target:
            with open(f"state_{step}_{loss}.mpack", 'wb') as file:
                   file.write(flax.serialization.to_bytes(driver.state))
        return True

``

[PhilipVinc]

Also notice that there is a logger that does exactly that https://netket.readthedocs.io/en/v3.6/api/_generated/driver/netket.logging.StateLog.html#netket.logging.StateLog