How do I get a regular pytorch model from BoTorchModel?

[cesare-montresor]

Hello there,

I'm new to BoTorch, I'm using it via Ax. I managed to run the optimization successfully, however, I ended up saving on the disk a BoTorchModel.
How do i get a regular PyTorch model out of a BoTorchModel ?

Thanks

[saitcakmak]

Hi @cesare-montresor. By "regular PyTorch model" are you referring to getting a BoTorch Model object as opposed to an instance of Ax's BoTorchModel? If that's the case, you can access the underlying BoTorch model instance using BoTorchModel.surrogate.model. Note that Ax does not save fitted BoTorch models to disc (it'll instead re-fit it as needed), so you'd have to save it manually.

[saitcakmak]

If you're using AxClient, the full path for the BoTorch model looks like this: ax_client.generation_strategy.model.model.surrogate.model (Ax overloads the word "model" quite a bit...). Note that this will only work if the model has been fitted already, which happens everytime a candidate is generated or can be manually triggered using ax_client.generation_strategy._fit_current_model(data=ax_client.experiment.lookup_data()).

[cesare-montresor]

Thanks for the really fast answer @saitcakmak !

Well, I'm still at my first experiment with Ax, I naively saved on the disk the model produced by the optimize function.
I now would like to use the trained model, but I realized that what I now have on the disk can't really be used directly.

This is what how I save it:

    best_parameters, values, experiment, model = optimize(
        parameters=params,
        evaluation_function=train_evaluate,
        objective_name='loss',
        minimize=True
    )
    means, covariances = values

    results = {
        'best_parameters': best_parameters,
        'means': means,
        'covariances': covariances,
    }
    print(results)
    session_dir = output_dir + session_id + '/'
    os.makedirs(session_dir, exist_ok=True)

    torch.save(model, session_dir + 'model.pth')
    results_yaml = yaml.dump(results) 
    with open(session_dir + 'params.yml', 'w') as f: f.write(results_yaml)

Then I try to load it for inference

    model = torch.load(session_dir+'model.pth')
    
    dataloader = getDataLoader(params, False)
    inference_model(model, dataloader)

[saitcakmak]

Since you're using torch.save/load, the model you have should be a ModelBridge object that has a trained BoTorch model attached to it, which you can access with model.model.surrogate.model. I don't know what the dataloader and inference_model are supposed to do, so I can't comment on those.

If you want to predict outcomes from the model, I'd actually recommend working with the Ax model directly. For this, you can call model.predict, defined here, to predict the outcomes for any given parameterization. The underlying BoTorch model lives in the transformed space, so working with it would require manually applying the Ax transforms to your parameters before passing them to the BoTorch model (and untransforming the model predictions).

[cesare-montresor]

Thank you very much, I'll try both methods to better understand how it works under the hood. Il lun 18 mar 2024, 21:17 Sait Cakmak ***@***.***> ha scritto:

…

Since you're using torch.save/load, the model you have should be a ModelBridge object that has a trained BoTorch model attached to it, which you can access with model.model.surrogate.model. I don't know what the dataloader and inference_model are supposed to do, so I can't comment on those. If you want to predict outcomes from the model, I'd actually recommend working with the Ax model directly. For this, you can call model.predict, defined here <https://github.com/facebook/Ax/blob/main/ax/modelbridge/base.py#L611>, to predict the outcomes for any given parameterization. The underlying BoTorch model lives in the transformed space, so working with it would require manually applying the Ax transforms to your parameters before passing them to the BoTorch model (and untransforming the model predictions). — Reply to this email directly, view it on GitHub <#2251 (reply in thread)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AADNK3OXVJLHWW5RNEDSK53YY5DVNAVCNFSM6AAAAABE4ESPUWVHI2DSMVQWIX3LMV43SRDJONRXK43TNFXW4Q3PNVWWK3TUHM4DQMZSG42DK> . You are receiving this because you were mentioned.Message ID: ***@***.***>

[cesare-montresor]

@saitcakmak I'm starting to think that I misunderstood what model refers to in the context of Ax and probably I didn't give you enough details on what I'm trying to achieve.

I'm using Ax to do find the ideal hyperparameters for my CNN, so what I'm really looking for is probably not BoTorchModel, but the model generated and fitted during the "best experiment", the one trained using the best parameters.
What is the correct naming for it? Is the experiment? or the candidate? Apologies for the sloppy questions, I'm still trying wrap my head around it.

For context, those are the parameters I'm tuning:

params = [
        {"name": "lr",                 "value_type": "float",  "type": "range", "bounds": [1e-6, 0.4], "log_scale": True}, #optimizer
        {"name": "lr_max",             "value_type": "float",  "type": "range", "bounds": [1e-6, 0.4], "log_scale": True}, #LR scheduler

        {"name": "batch_size",         "value_type": "int",    "type": "range", "bounds": [4, 128]}, #train loop
        {"name": "num_epoch",          "value_type": "int",    "type": "range", "bounds": [1, 30]}, #train loop
        {"name": "drop_out",           "value_type": "float",  "type": "range", "bounds": [0.0, 0.9]}, # model FC
        {"name": "fc_hidden_num",      "value_type": "int",    "type": "range", "bounds": [0, 10]}, # model FC
        {"name": "fc_hidden_size",     "value_type": "int",    "type": "range", "bounds": [64, 2048]}, # model FC
        
        {"name": "resnet_size",        "value_type": "int",    "type": "choice", "values": [18, 34, 50, 101, 152], "sort_values": True, "is_ordered": True}, # model CNN
    ]

[cesare-montresor]

Nevermind, I realized i can dump any model withtin the train_eval function, also, given the best parameters, i can run one more training using the best_paramters and save the tuned model.
I was getting lost in a glass of water, thank for the support!