Update load_state_dict for ModelList to support fully Bayesian models (#1395)

Summary:
Pull Request resolved: #1395

Update `ModelList` to initialize `SaasFullyBayesianSingleTaskGP` before loading the `state_dict`. This is needed in order to do cross-validation in MBM.

Reviewed By: Balandat, qingfeng10

Differential Revision: D39402322

fbshipit-source-id: a3cb8144585cfd635178d4918084d7391f985262

Author: dme65

botorch/models/model.py

@@ -16,7 +16,7 @@
 from abc import ABC, abstractmethod
 from collections import defaultdict
 from copy import deepcopy
-from typing import Any, Callable, Dict, Hashable, List, Optional, Union
+from typing import Any, Callable, Dict, Hashable, List, Mapping, Optional, Union
 
 import numpy as np
 import torch
@@ -420,3 +420,17 @@ def transform_inputs(self, X: Tensor) -> List[Tensor]:
             except AttributeError:
                 transformed_X_list.append(X)
         return transformed_X_list
+
+    def load_state_dict(
+        self, state_dict: Mapping[str, Any], strict: bool = True
+    ) -> None:
+        """Initialize the fully Bayesian models before loading the state dict."""
+        for i, m in enumerate(self.models):
+            if is_fully_bayesian(m):
+                filtered_dict = {
+                    k.replace(f"models.{i}.", ""): v
+                    for k, v in state_dict.items()
+                    if k.startswith(f"models.{i}")
+                }
+                m.load_state_dict(filtered_dict)
+        super().load_state_dict(state_dict=state_dict, strict=strict)

test/models/test_fully_bayesian.py

@@ -326,23 +326,26 @@ def test_fit_model(self):
             self.assertEqual(median_lengthscale.shape, torch.Size([4]))
             self.assertEqual(model.num_mcmc_samples, 3)
 
-            # Test loading via state dict
-            state_dict = model.state_dict()
+            # Check the keys in the state dict
             true_keys = EXPECTED_KEYS_NOISE if infer_noise else EXPECTED_KEYS
-            self.assertEqual(set(state_dict.keys()), set(true_keys))
-            _, _, _, model_new = self._get_data_and_model(
-                infer_noise=infer_noise, **tkwargs
-            )
-            self.assertEqual(model_new.state_dict(), {})
-            model_new.load_state_dict(state_dict)
-            self.assertEqual(model.state_dict().keys(), model_new.state_dict().keys())
-            for k in model.state_dict().keys():
-                self.assertTrue(
-                    (model.state_dict()[k] == model_new.state_dict()[k]).all()
+            self.assertEqual(set(model.state_dict().keys()), set(true_keys))
+
+            for i in range(2):  # Test loading via state dict
+                m = model if i == 0 else ModelList(model, deterministic)
+                state_dict = m.state_dict()
+                _, _, _, m_new = self._get_data_and_model(
+                    infer_noise=infer_noise, **tkwargs
                 )
-            preds1, preds2 = model.posterior(test_X), model_new.posterior(test_X)
-            self.assertTrue((preds1.mean == preds2.mean).all())
-            self.assertTrue((preds1.variance == preds2.variance).all())
+                m_new = m_new if i == 0 else ModelList(m_new, deterministic)
+                if i == 0:
+                    self.assertEqual(m_new.state_dict(), {})
+                m_new.load_state_dict(state_dict)
+                self.assertEqual(m.state_dict().keys(), m_new.state_dict().keys())
+                for k in m.state_dict().keys():
+                    self.assertTrue((m.state_dict()[k] == m_new.state_dict()[k]).all())
+                preds1, preds2 = m.posterior(test_X), m_new.posterior(test_X)
+                self.assertTrue(torch.equal(preds1.mean, preds2.mean))
+                self.assertTrue(torch.equal(preds1.variance, preds2.variance))
 
             # Make sure the model shapes are set correctly
             self.assertEqual(model.pyro_model.train_X.shape, torch.Size([n, d]))