Merge branch 'master' into rename_deep_river

pyproject.toml

@@ -49,6 +49,7 @@ module = [
     "pandas.*",
     "scipy.*",
     "torch.*",
+    "torchviz.*",
     "graphviz.*",
     "vaex.*",
     "torch.*",

river_torch/anomaly/rolling_ae.py

@@ -219,7 +219,8 @@ def score_one(self, x: dict) -> float:
             x_win.append(list(x.values()))
             x_t = deque2rolling_tensor(x_win, device=self.device)
             self.module.eval()
-            x_pred = self.module(x_t)
+            with torch.inference_mode():
+                x_pred = self.module(x_t)
             loss = self.loss_fn(x_pred, x_t)
             res = loss.item()
 
@@ -240,7 +241,8 @@ def score_many(self, X: pd.DataFrame) -> List[Any]:
         if len(self._x_window) == self.window_size:
             self.module.eval()
             X_t = deque2rolling_tensor(x_win, device=self.device)
-            x_pred = self.module(X_t)
+            with torch.inference_mode():
+                x_pred = self.module(X_t)
             loss = torch.mean(
                 self.loss_fn(x_pred, x_pred, reduction="none"),
                 dim=list(range(1, x_pred.dim())),

river_torch/base.py

@@ -8,6 +8,14 @@
 
 from river_torch.utils import get_loss_fn, get_optim_fn
 
+try:
+    from graphviz import Digraph
+    from torchviz import make_dot
+except ImportError as e:
+    raise ValueError(
+        "You have to install graphviz to use the draw method"
+    ) from e
+
 
 class DeepEstimator(base.Estimator):
     """
@@ -106,6 +114,15 @@ def _filter_kwargs(self, fn: Callable, override=None, **kwargs) -> dict:
         res.update(override)
         return res
 
+    def draw(self) -> Digraph:
+        """Draws the wrapped model."""
+        first_parameter = next(self.module.parameters())
+        input_shape = first_parameter.size()
+        y_pred = self.module(torch.rand(input_shape))
+        return make_dot(
+            y_pred.mean(), params=dict(self.module.named_parameters())
+        )
+
     def initialize_module(self, **kwargs):
         """
         Parameters

river_torch/classification/classifier.py

@@ -254,7 +254,8 @@ def predict_proba_one(self, x: dict) -> Dict[ClfTarget, float]:
             self.initialize_module(**self.kwargs)
         x_t = dict2tensor(x, device=self.device)
         self.module.eval()
-        y_pred = self.module(x_t)
+        with torch.inference_mode():
+            y_pred = self.module(x_t)
         return output2proba(
             y_pred, self.observed_classes, self.output_is_logit
         )
@@ -306,7 +307,8 @@ def predict_proba_many(self, X: pd.DataFrame) -> pd.DataFrame:
             self.initialize_module(**self.kwargs)
         X_t = df2tensor(X, device=self.device)
         self.module.eval()
-        y_preds = self.module(X_t)
+        with torch.inference_mode():
+            y_preds = self.module(X_t)
         return pd.Dataframe(output2proba(y_preds, self.observed_classes))
 
     def _adapt_output_dim(self):

river_torch/classification/rolling_classifier.py

@@ -254,12 +254,12 @@ def predict_proba_one(self, x: dict) -> Dict[ClfTarget, float]:
             self.initialize_module(**self.kwargs)
 
         if len(self._x_window) == self.window_size:
-            self.module.eval()
-            x_win = self._x_window.copy()
-            x_win.append(list(x.values()))
-            x_t = deque2rolling_tensor(x_win, device=self.device)
-            y_pred = self.module(x_t)
-            proba = output2proba(y_pred, self.observed_classes)
+            with torch.inference_mode():
+                x_win = self._x_window.copy()
+                x_win.append(list(x.values()))
+                x_t = deque2rolling_tensor(x_win, device=self.device)
+                y_pred = self.module(x_t)
+                proba = output2proba(y_pred, self.observed_classes)
         else:
             proba = self._get_default_proba()
 
@@ -322,12 +322,12 @@ def predict_proba_many(self, X: pd.DataFrame) -> pd.DataFrame:
         x_win.extend(X.values.tolist())
 
         if len(x_win) == self.window_size:
-            self.module.eval()
-            x_t = deque2rolling_tensor(x_win, device=self.device)
-            probas = self.module(x_t).detach().tolist()
-            if len(probas) < len(X):
-                default_proba = self._get_default_proba()
-                probas = [default_proba] * (len(X) - len(probas)) + probas
+            with torch.inference_mode():
+                x_t = deque2rolling_tensor(x_win, device=self.device)
+                probas = self.module(x_t).detach().tolist()
+                if len(probas) < len(X):
+                    default_proba = self._get_default_proba()
+                    probas = [default_proba] * (len(X) - len(probas)) + probas
         else:
             default_proba = self._get_default_proba()
             probas = [default_proba] * len(X)

river_torch/regression/regressor.py

@@ -177,7 +177,9 @@ def predict_one(self, x: dict) -> RegTarget:
             self.initialize_module(**self.kwargs)
         x_t = dict2tensor(x, self.device)
         self.module.eval()
-        return self.module(x_t).item()
+        with torch.inference_mode():
+            y_pred = self.module(x_t).item()
+        return y_pred
 
     def learn_many(self, X: pd.DataFrame, y: List) -> "Regressor":
         """
@@ -223,4 +225,6 @@ def predict_many(self, X: pd.DataFrame) -> List:
 
         X = df2tensor(X, device=self.device)
         self.module.eval()
-        return self.module(X).detach().tolist()
+        with torch.inference_mode():
+            y_preds = self.module(X).detach().tolist()
+        return y_preds

river_torch/regression/rolling_regressor.py

@@ -192,11 +192,11 @@ def predict_one(self, x: dict) -> RegTarget:
             self.initialize_module(**self.kwargs)
 
         if len(self._x_window) == self.window_size:
-            self.module.eval()
-            x_win = self._x_window.copy()
-            x_win.append(list(x.values()))
-            x_t = deque2rolling_tensor(x_win, device=self.device)
-            res = self.module(x_t).detach().numpy().item()
+            with torch.inference_mode():
+                x_win = self._x_window.copy()
+                x_win.append(list(x.values()))
+                x_t = deque2rolling_tensor(x_win, device=self.device)
+                res = self.module(x_t).detach().numpy().item()
 
         if self.append_predict:
             self._x_window.append(list(x.values()))
@@ -212,9 +212,9 @@ def predict_many(self, X: pd.DataFrame) -> List:
         x_win = self._x_window.copy()
         x_win.extend(X.values.tolist())
         if len(x_win) == self.window_size:
-            self.module.eval()
-            x_t = deque2rolling_tensor(x_win, device=self.device)
-            res = self.module(x_t).detach().tolist()
+            with torch.inference_mode():
+                x_t = deque2rolling_tensor(x_win, device=self.device)
+                res = self.module(x_t).detach().tolist()
 
         if self.append_predict:
             self._x_window.extend(X.values.tolist())

river_torch/utils/hooks.py

@@ -12,8 +12,9 @@ def __call__(self, module, input, output):
             self.ordered_modules.append(module)
 
 
-def apply_hooks(module, hook, handles):
+def apply_hooks(module, hook, handles=[]):
     for child in module.children():
         apply_hooks(child, hook, handles)
     handle = module.register_forward_hook(hook)
     handles.append(handle)
+    return handles

setup.py

@@ -16,10 +16,11 @@
     "scikit-learn~=1.0.2",
     "torch~=1.13.0",
     "pandas~=1.3.2",
-    "numpy~=1.23.2",
+    "numpy~=1.24.0",
     "river~=0.14.0",
     "tqdm~=4.61.2",
     "ordered-set~=4.1.0",
+    "torchviz~=0.0.2",
 ]
 
 dev_packages = base_packages + [