surrogates_interface 4.1.0__tar.gz → 4.1.1__tar.gz

{surrogates_interface-4.1.0 → surrogates_interface-4.1.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: surrogates_interface
-Version: 4.1.0
+Version: 4.1.1
 Summary: Interface for surrogate models.
 Home-page: https://gitlab.windenergy.dtu.dk/surrogate-models/surrogate-models-interface
 License: MIT

{surrogates_interface-4.1.0 → surrogates_interface-4.1.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "surrogates_interface"
-version = "4.1.0"
+version = "4.1.1"
 description = "Interface for surrogate models."
 authors = ["DTU Wind and Energy Systems"]
 maintainers = [

{surrogates_interface-4.1.0 → surrogates_interface-4.1.1}/surrogates_interface/surrogates.py

@@ -1030,48 +1030,51 @@ class SurrogateModelFamily:
 
 
 class PyTorchModel(SurrogateModel):
-    import torch
-    import torch._dynamo  # fmt: skip
-    torch._dynamo.config.suppress_errors = True
-    import torch.func as torch_func  # fmt: skip
+    """
+    Class that stores and evaluates an existing PyTorch model.
+    """
 
-    # from torch.func import jacfwd, jacrev, jvp, vjp, vmap  # fmt: skip
+    def __init__(self, *args, dtype=None, max_batch=50_000, **kwargs):
+        import torch
+        import torch._dynamo
 
-    def __init__(self, *args, dtype=torch.float16, max_batch=50_000, **kwargs):
+        torch._dynamo.config.suppress_errors = True
         super().__init__(*args, **kwargs)
         self.model.eval()
-        self.device = self.torch.device(
-            "cuda" if self.torch.cuda.is_available() else "cpu"
-        )
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         self.model.to(self.device)
-        self.dtype = dtype
-        self.model.to(dtype)
+        self.dtype = torch.float16 if dtype is None else dtype
+        self.model.to(self.dtype)
         self.max_batch = max_batch
 
     def __to_numpy(self, _x):
+        import torch
+
         try:
             import cupy as cp
         except ImportError:  # pragma: no cover
             cp = np
-        lalib = (
-            cp if self.device == self.torch.device("cuda") else np
-        )  # pragma: no cover
+        lalib = cp if self.device == torch.device("cuda") else np  # pragma: no cover
         return lalib.asarray(_x).astype(lalib.float32)
 
     def __to_tensor(self, _x):
+        import torch
+
         if len(_x.shape) < 2:
             _x = _x.reshape(1, -1)
-        return self.torch.as_tensor(_x, device=self.device, dtype=self.dtype)
+        return torch.as_tensor(_x, device=self.device, dtype=self.dtype)
 
     def __record_in_dtype(self, _x):  # pragma: no cover
         # TODO: could be used to retain the input/output dtype the same
         self.in_dtype = _x.dtype
 
     def _predict_model_output(self, x):
+        import torch
+
         x = self.__to_tensor(x)
 
-        with self.torch.no_grad():
-            output = self.torch.zeros(
+        with torch.no_grad():
+            output = torch.zeros(
                 (x.shape[0], self.n_outputs),
                 device=self.device,
                 dtype=self.dtype,
@@ -1100,21 +1103,28 @@ class PyTorchModel(SurrogateModel):
         return self.__to_numpy(y)
 
     def _save_model(self, filepath):
-        self.torch.save(self.model.state_dict(), filepath)
+        import torch
+
+        torch.save(self.model.state_dict(), filepath)
 
     @staticmethod
     def _load_model(filepath, model_cls):
-        import torch  # fmt: skip
+        import torch
+
         model_cls.load_state_dict(torch.load(filepath, weights_only=True))
         return model_cls
 
     def _predict_model_jvp(self, x, dx):
+        import torch.func as torch_func
+
         if not hasattr(self, "batched_jvp_func"):
             def g(_x):  # fmt: skip
                 return self.model(_x)
-            def single_jvp(_x, _v):  # fmt: skip
-                return self.torch_func.jvp(g, (_x,), (_v,))
-            self.batched_jvp_func = self.torch_func.vmap(single_jvp)
+
+            def single_jvp(_x, _v):
+                return torch_func.jvp(g, (_x,), (_v,))
+
+            self.batched_jvp_func = torch_func.vmap(single_jvp)
 
         x = self.__to_tensor(x)
         dx = self.__to_tensor(dx)
@@ -1122,13 +1132,17 @@ class PyTorchModel(SurrogateModel):
         return self.__to_numpy(output.detach()), self.__to_numpy(jvp_res.detach())
 
     def _predict_model_vjp(self, x, dy):
+        import torch.func as torch_func
+
         if not hasattr(self, "batched_vjp_func"):
             def g(_x):  # fmt: skip
                 return self.model(_x)
-            def single_vjp(_x, _v):  # fmt: skip
-                out, f = self.torch_func.vjp(g, _x)
+
+            def single_vjp(_x, _v):
+                out, f = torch_func.vjp(g, _x)
                 return out, f(_v)
-            self.batched_vjp_func = self.torch_func.vmap(single_vjp)
+
+            self.batched_vjp_func = torch_func.vmap(single_vjp)
 
         x = self.__to_tensor(x)
         dy = self.__to_tensor(dy)
@@ -1136,49 +1150,50 @@ class PyTorchModel(SurrogateModel):
         return self.__to_numpy(output.detach()), self.__to_numpy(vjp_res[0].detach())
 
     def _predict_model_output_and_jacobian_fwd(self, x):
+        import torch.func as torch_func
+
         if not hasattr(self, "y_jac_fwd_func"):
             def g(_x):  # fmt: skip
                 result = self.model(_x)
                 return (result, result)
-            self.y_jac_fwd_func = self.torch_func.vmap(
-                self.torch_func.jacfwd(g, has_aux=True)
-            )
+            self.y_jac_fwd_func = torch_func.vmap(torch_func.jacfwd(g, has_aux=True))
 
         x = self.__to_tensor(x)
         jac, output = self.y_jac_fwd_func(x)
         return self.__to_numpy(output.detach()), self.__to_numpy(jac.detach())
 
     def _predict_model_output_and_jacobian_rev(self, x):
+        import torch.func as torch_func
+
         if not hasattr(self, "y_jac_rev_func"):
             def g(_x):  # fmt: skip
                 result = self.model(_x)
                 return (result, result)
-            self.y_jac_rev_func = self.torch_func.vmap(
-                self.torch_func.jacrev(g, has_aux=True)
-            )
+            self.y_jac_rev_func = torch_func.vmap(torch_func.jacrev(g, has_aux=True))
 
         x = self.__to_tensor(x)
         jac, output = self.y_jac_rev_func(x)
         return self.__to_numpy(output.detach()), self.__to_numpy(jac.detach())
 
     def _predict_model_jacobian_sbatch(self, x):  # pragma: no cover
-        """Not used for now, but left as a reference for future if GPU memory issues arise."""
+        # Not used for now, but left as a reference for future if GPU memory issues arise.
+        import torch
+        import torch.func as torch_func
+
         if not hasattr(self, "y_jac_func"):
             def g(_x):  # fmt: skip
                 result = self.model(_x)
                 return (result, result)
-            self.y_jac_func = self.torch_func.vmap(
-                self.torch_func.jacrev(g, has_aux=True)
-            )
+            self.y_jac_func = torch_func.vmap(torch_func.jacrev(g, has_aux=True))
 
         x = self.__to_tensor(x)
         batch_size = x.shape[0]
-        y = self.torch.zeros(
+        y = torch.zeros(
             (batch_size, self.n_outputs),
             device=self.device,
             dtype=self.dtype,
         )
-        jac = self.torch.zeros(
+        jac = torch.zeros(
             (batch_size, self.n_outputs, self.n_inputs),
             device=self.device,
             dtype=self.dtype,