Patch state_dict loading for PairwiseGP (#1359)

Summary:
Pull Request resolved: #1359

Ensure custom handling gets applied even when `load_state_dict` is called on a parent module.

Reviewed By: saitcakmak

Differential Revision: D38881840

fbshipit-source-id: 8ce2e6026a8f7359d829f79a04ea1e543703a4c4

Author: James Wilson

botorch/models/pairwise_gp.py

@@ -27,6 +27,7 @@
 import numpy as np
 import torch
 from botorch.acquisition.objective import PosteriorTransform
+from botorch.exceptions import UnsupportedError
 from botorch.models.likelihoods.pairwise import (
     PairwiseLikelihood,
     PairwiseProbitLikelihood,
@@ -712,35 +713,49 @@ def set_train_data(
         self.to(self.datapoints)
 
     def load_state_dict(
-        self, state_dict: Dict[str, Tensor], strict: Optional[bool] = False
+        self, state_dict: Dict[str, Tensor], strict: bool = False
     ) -> _IncompatibleKeys:
         r"""Removes data related buffers from the `state_dict` and calls
         `super().load_state_dict` with `strict=False`.
 
         Args:
             state_dict: The state dict.
-            strict: A boolean denoting whether to error out if all keys are not
-                present in the `state_dict`. Since we remove data related buffers
-                from the `state_dict`, this will lead to an error whenever
-                `strict=True`. Instead, we overwrite it with `strict=False`, and
-                raise a warning explaining this if `strict=True` is passed.
+            strict: Boolean specifying whether or not given and instance-bound
+                state_dicts should have identical keys. Only implemented for
+                `strict=False` since buffers will filters out when calling
+                `_load_from_state_dict`.
 
         Returns:
             A named tuple `_IncompatibleKeys`, containing the `missing_keys`
-            and `unexpected_keys`. Note that the buffers we remove from the
-            `state_dict` may be listed under `missing_keys`.
+            and `unexpected_keys`.
         """
         if strict:
-            warnings.warn(
-                f"Received `strict=True` in {self.__class__.__name__}.load_state_dict "
-                "call. This will be overwritten with `strict=False` after removing "
-                "a set of data related buffers from the `state_dict`.",
-                RuntimeWarning,
-            )
-        for key in self._buffer_names:
-            state_dict.pop(key, None)
+            raise UnsupportedError("Passing strict=True is not supported.")
+
         return super().load_state_dict(state_dict=state_dict, strict=False)
 
+    def _load_from_state_dict(
+        self,
+        state_dict: Dict[str, Tensor],
+        prefix: str,
+        local_metadata: Dict[str, Any],
+        strict: bool,
+        missing_keys: List[str],
+        unexpected_keys: List[str],
+        error_msgs: List[str],
+    ) -> None:
+        super()._load_from_state_dict(
+            state_dict={
+                k: v for k, v in state_dict.items() if k not in self._buffer_names
+            },
+            prefix=prefix,
+            local_metadata=local_metadata,
+            strict=False,
+            missing_keys=missing_keys,
+            unexpected_keys=unexpected_keys,
+            error_msgs=error_msgs,
+        )
+
     def forward(self, datapoints: Tensor) -> MultivariateNormal:
         r"""Calculate a posterior or prior prediction.
 

test/models/test_pairwise_gp.py

@@ -10,7 +10,7 @@
 import torch
 from botorch import fit_gpytorch_model
 from botorch.acquisition.objective import ScalarizedPosteriorTransform
-from botorch.exceptions.warnings import OptimizationWarning
+from botorch.exceptions import OptimizationWarning, UnsupportedError
 from botorch.models.likelihoods.pairwise import (
     PairwiseLogitLikelihood,
     PairwiseProbitLikelihood,
@@ -320,11 +320,14 @@ def test_fantasize(self):
     def test_load_state_dict(self):
         model, _ = self._get_model_and_data(batch_shape=[])
         sd = model.state_dict()
-        with warnings.catch_warnings(record=True) as ws:
-            missing, unexpected = model.load_state_dict(sd, strict=True)
-        # Check that the warning was raised.
-        self.assertTrue(any("strict=True" in str(w.message) for w in ws))
-        # Check that buffers are missing.
-        self.assertIn("datapoints", missing)
-        self.assertIn("D", missing)
-        self.assertIn("covar", missing)
+        with self.assertRaises(UnsupportedError):
+            model.load_state_dict(sd, strict=True)
+
+        # Set instance buffers to None
+        for buffer_name in model._buffer_names:
+            model.register_buffer(buffer_name, None)
+
+        # Check that instance buffers were not restored
+        _ = model.load_state_dict(sd)
+        for buffer_name in model._buffer_names:
+            self.assertIsNone(model.get_buffer(buffer_name))