Restart training from a saved file

[zhyang-dev]

Dear Miles,

I have been training a large number of models in parallel and have achieved some relatively good results. However, they did not quite meet my expectations and I am interested in further training. I am wondering if there is any feature that supports continuing training from a selected checkpoint.

Thank you for your attention to my question.

[MilesCranmer]

You can use the warm_start=True parameter to continue where you left off.

However after exiting Python, it is not possible to restart, as the Julia runtime will be closed and relevant variables erased. Perhaps one could serialize the Julia variables (model.raw_julia_state_) but I have not tried this.

Edit: note that PySR v1.x saves the search state to a file, so this point above doesn't apply.

[zhyang-dev]

Thank you for your response. I tried loading the model from a file and setting the warm_start parameter to True before starting the training. The training process will restart. As you mentioned, this may be due to the inability to recover Julia variables.
I was thinking that perhaps there could be an option to determine whether to serialize Julia's context for continued training. This might be helpful for those who need to restart the training process.
For now, i will try to run the trainning forever and manually restart some of the slow loss reducing cases. Again, thank you for your help.

[MilesCranmer]

Right: you can’t have a successful warm start without raw_julia_state_, and that is cleared from the state when the PySRRegressor is saved to a pickle file.

However you could try to save this object using the Serialization library in Julia: https://docs.julialang.org/en/v1/stdlib/Serialization/. You will need to access this with PyJulia,
for example:

from julia import Serialization

Serialization.serialize(
    "checkpoint.pt",
    model.raw_julia_state_  # This is the trained model
)

Then, in a new process:

from pysr import PySRRegressor

model = PySRRegressor.from_file("hall_of_fame...pkl")
model.warm_start = True

from pysr.julia_helpers import init_julia

init_julia()

from julia import SymbolicRegression  # Needed to load library (usually this is done by .fit())
from julia import Serialization

model.raw_julia_state_ = Serialization.deserialize("checkpoint.pt")

Now it should be an identical model as before you closed Python.

[MilesCranmer]

@zhyang-dev would you be up for making a PR that adds this ability to __setstate__ and __getstate__ of PySRRegressor?

[MilesCranmer]

@tttc3 you might be interested in this too. It looks like you can fully save the entire model state, even containing the Julia variables. i.e., you could restart a PySR search in a new Python runtime.

[Pulpas]

Dear Miles, tttc3, and Zhyang,

I am also training several large number of models in parallel and have achieved some relatively good results.
I would also like to restart from the previous model, which I thought was possible using the function:

model = PySRRegressor.from_file(str("file-name.pkl"), warm_start=True)

combined with the option warm_start set to True, to start where the model was left off.

However, it seems that the model does not start from it was left off.

So I looked it up and reached this thread where is it prescribed to save the Julia state, by doing

from julia import Serialization

Serialization.serialize(
    "checkpoint.pt",
    model.raw_julia_state_
)

which I call after the model.fit() call.

Then, I restart from the Julia state by loading the checkpoint.pt file along with the .pkl file doing

        model = PySRRegressor.from_file(str("file-namepkl"), warm_start=True)
        
        from pysr.julia_helpers import init_julia

        init_julia()

        from julia import SymbolicRegression  
        from julia import Serialization

        model.raw_julia_state_ = Serialization.deserialize("checkpoint.pt")
        
        model.fit(X, y)

However, this does not restart from the saved state as I thought it would.

Is there anything I'm missing out ? Did @tttc3, or @zhyang-dev successfully restart from previous state ?

Best,

[Pulpas]

Dear @MilesCranmer,

The results of the function call:

from julia import Serialization

Serialization.serialize(
    "checkpoint.pt",
    model.raw_julia_state_
)

generates a non empty checkpoint.pt file. To make sure that model is not empty, I started a new python instance where I print the saved model before and after the function call.

I agree with you, that would be a very interesting feature to have in PySR. My PI keeps me very busy right now, but I'll try to propose an implementation for this feature in the near future.

[MilesCranmer]

Did you try setting the warm start manually, after initializing it?

[Pulpas]

Your following suggestion,

model = PySRRegressor.from_file("hall_of_fame...pkl")
model.warm_start = True

enable loading the previous model ! Thanks a lot for pointing this out !

[MilesCranmer]

Awesome! We should really have kwargs passed automatically through from_file

[Pulpas]

I totally agree, I was a bit naive on that one, I should have checked the source code!

[lhabersbrunner]

[MilesCranmer]

The suggestions in this thread were for PySR 0.x. On PySR 1.x, this should happen automatically when you load from a checkpoint. Is there a particular reason you are trying to set it manually?

(On PySR 1.x, the attribute is julia_state_stream_: NDArray[np.uint8] | None. Note that this stores the serialized data directly - a binary stream of data. The deserialization automatically happens now, whenever you call model.julia_state_. But the state stream is automatically loaded from a checkpoint file, so you shouldn't need to worry about this. But perhaps it's not working for you?)

[lhabersbrunner]

Thank you for the quick response! I am having problems with the 1.X version method.

When using the "model = PySRRegressor.from_file(run_directory=os.path.dirname(pretrained_model_path_pkl))" command the best equation does not change no matter how I change the iterations with model.niterations after calling the above mentioned from_file function.

This is how I have tried it so far:

model = PySRRegressor.from_file(run_directory=os.path.dirname(pretrained_model_path_pkl))
model.niterations = niterations
model.maxsize = maxsize
model.warm_start = True
model.verbosity = 1
model.fit(X_train, y_train, variable_names=[f"{var_short}"])

As this hasn't worked so far, I wanted to try it manually.

[MilesCranmer]

Can you make an issue with the full code so I can reproduce it?

[lhabersbrunner]

Please excuse my late response.
I ran the original code in fiels from a OneDrive folder and rerun with the same files locally on my PC and this seemed to fix the problem. I assume that the link to the OneDrive folder was the problem.

[gm89uk]

In the newest update for symbolicregression.jl it's really easy to load an old model:

using DataFrames, CSV
guesses = CSV.read("<dir>\\QuickLoad\\hall_of_fame.csv", DataFrame).Equation #hall_of_fame.csv you want to load
#...
model = SRRegressor(
    #... 
    guesses = guesses,
    #...  
)
mach = machine(model, X, y)
fit!(mach)