Misc Fixes (#14)

t-vi · web-flow · commit d828e41247de · 2022-06-09T21:29:47.000+02:00
- Python Optimal Transport now wants numpy arrays and does not
  like tensors.
- Saving and Loading when buffers are None...
diff --git a/test/test_reducers.py b/test/test_reducers.py
@@ -17,6 +17,13 @@ def test_pca():
     red2 = pca(b)
     red2_ref = torch.from_numpy(pca_ref.transform(b))
 
+def test_reducer_load_save():
+    pca = torchdrift.reducers.PCAReducer(n_components=2)
+    a = torch.randn(100, 50, dtype=torch.double)
+    red = pca.fit(a)
+    pca2 = torchdrift.reducers.PCAReducer(n_components=2)
+    pca2.load_state_dict(pca.state_dict())
+    red2 = pca2(a)
 
 def test_reducer():
     x = torch.randn(5, 5)
diff --git a/torchdrift/detectors/detector.py b/torchdrift/detectors/detector.py
@@ -63,3 +63,11 @@ def forward(
         if self.return_p_value:
             return p_value
         return ood_score
+
+    def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict,
+                              missing_keys, unexpected_keys, error_msgs):
+        for bname, b in self._buffers.items():
+            if prefix + bname in state_dict and b is None:
+                setattr(self, bname, state_dict[prefix + bname])
+        super()._load_from_state_dict(state_dict, prefix, local_metadata, strict,
+                                      missing_keys, unexpected_keys, error_msgs)
diff --git a/torchdrift/detectors/wasserstein.py b/torchdrift/detectors/wasserstein.py
@@ -47,7 +47,7 @@ def wasserstein(x, y, p=2.0, fraction_to_match=1.0, n_perm=1000, return_coupling
     if fraction_to_match < 1.0:
         weights_y[:, -1] = 1.0 - fraction_to_match
 
-    coupling = torch.from_numpy(ot_emd(weights_x[0], weights_y[0], dists_p.cpu()))
+    coupling = torch.from_numpy(ot_emd(weights_x[0].numpy(), weights_y[0].numpy(), dists_p.cpu().numpy()))
 
     if (coupling[:, :num_y].sum() / fraction_to_match - 1).abs().item() > 1e-5:  # pragma: no cover
         raise RuntimeError("Numerical stability failed")
diff --git a/torchdrift/reducers/pca.py b/torchdrift/reducers/pca.py
@@ -16,6 +16,8 @@ class PCAReducer(Reducer):
     def __init__(self, n_components: int = 2):
         super().__init__()
         self.n_components = n_components
+        self.register_buffer("mean", None)
+        self.register_buffer("comp", None)
 
     def extra_repr(self) -> str:
         return f"n_components={self.n_components}"
@@ -37,3 +39,11 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         x = x - self.mean
         reduced = x @ self.comp
         return reduced
+
+    def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict,
+                              missing_keys, unexpected_keys, error_msgs):
+        for bname in ("mean", "comp"):
+            if prefix + bname in state_dict:
+                setattr(self, bname, state_dict[prefix + bname])
+        super()._load_from_state_dict(state_dict, prefix, local_metadata, strict,
+                                      missing_keys, unexpected_keys, error_msgs)
diff --git a/torchdrift/reducers/reducer.py b/torchdrift/reducers/reducer.py
@@ -19,3 +19,11 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
         Override this in your reducer implementation."""
         raise NotImplementedError("Override forward in subclass")
+
+    def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict,
+                              missing_keys, unexpected_keys, error_msgs):
+        for bname, b in self._buffers.items():
+            if prefix + bname in state_dict and b is None:
+                setattr(self, bname, state_dict[prefix + bname])
+        super()._load_from_state_dict(state_dict, prefix, local_metadata, strict,
+                                      missing_keys, unexpected_keys, error_msgs)
