Saving and loading Botorch models

[danieljacobsen-ai]

I'm trying to store and retrieve a Botorch model. Specifically a SaasFullyBayesianSingleTaskGP.

I'm using the Gpytorch documentation to store and later load the state_dict into the model.

Saving:

torch.save(gp.state_dict(), 'model_state.pth')
pickle.dump([train_X, train_Y], open('train_data.pkl', 'wb'))

Loading:

train_X, train_Y = pickle.load(open('train_data.pkl', 'rb'))
state_dict = torch.load('data', 'model_state.pth')
gp = SaasFullyBayesianSingleTaskGP(train_X=train_X, train_Y=train_Y)
gp.load_state_dict(state_dict)

print(gp.state_dict())

>>OrderedDict()

That is, the state_dict is empty. Upon saving, the state_dict contains the trained parameters as expected.

The data state_dict is identical - but it seem to fail to load into the model.

I'm obviously doing something wrong. What is an effective way of storing and retrieving such models?

Thanks,
Daniel

[Balandat]

I'm not sure whether we've tested saving/loading of the fully Bayesian models. They're doing a few things differently than standard models that could potentially be an issue here (e.g. first of all that the underlying PyroModel used internally doesn't subclass torch.nn.Module. @dme65, any thoughts here?

[dme65]

Nice catch, this hasn't been tested and currently isn't supported since SaasFullyBayesianSingleTaskGP initializes covar_module, mean_module, and likelihood to None. This means that the necessary modules don't exist when you call load_state_dict which explains why the state_dict ends up being empty. Supporting this will likely require PyroModel to initialize covar_module, mean_module, and likelihood when the SaasFullyBayesianSingleTaskGP model is created.

[Balandat]

@dme65 how much work will it be to put up a PR to support this?

[dme65]

@Balandat Hopefully not too bad. The main thing that will be tricky is to deal properly with the batch_shape. We don't know the number of MCMC samples by the time the model is created, so we may have to modify some of the load hooks to reshape the parameters to match the state dict. Let me know you if you have any thoughts here.

[dme65]

Sorry for the delay here! This was fixed in #1384.