Added option of multi tabular components. Runs with model fusion. More tests of course needed

Author: jrzaurin

pytorch_widedeep/models/model_fusion.py

@@ -110,6 +110,7 @@ def __init__(
                 "max",
                 "sum",
                 "mult",
+                "dot",
                 "head",
             ],
             List[Literal["concatenate", "mean", "max", "sum", "mult", "head"]],
@@ -182,11 +183,20 @@ def forward(self, X: List[Tensor]) -> Tensor:  # noqa: C901
             return self.head(
                 torch.cat([model(x) for model, x in zip(self.models, X)], -1)
             )
+        elif self.fusion_method == "dot":
+            assert len(X) == 2, (
+                "When using 'dot' as fusion_method, only two models "
+                " can be fused. Accordingly, only two inputs should be provided"
+            )
+            outputs = [model(x) for model, x in zip(self.models, X)]
+            return torch.bmm(outputs[1].unsqueeze(1), outputs[0].unsqueeze(2)).view(
+                -1, 1
+            )
         else:
             if isinstance(self.fusion_method, str):
                 fusion_methods = [self.fusion_method]
             else:
-                fusion_methods = self.fusion_method
+                fusion_methods = self.fusion_method  # type: ignore
 
             fused_outputs: List[Tensor] = []
             for fm in fusion_methods:
@@ -210,7 +220,7 @@ def forward(self, X: List[Tensor]) -> Tensor:  # noqa: C901
                     else:
                         # This should never happen, but avoids type errors
                         raise ValueError(
-                            "fusion_method must be one of ['concatenate', 'mean', 'max', 'sum', 'mult', 'head'] "
+                            "fusion_method must be one of ['concatenate', 'mean', 'max', 'sum', 'mult', 'dot', 'head'] "
                             "or a list of those"
                         )
                 fused_outputs.append(out)
@@ -260,12 +270,14 @@ def output_dim(self) -> int:
                 if hasattr(self, "head_hidden_dims")
                 else self.head.output_dim
             )
+        elif self.fusion_method == "dot":
+            output_dim = 1
         else:
             output_dim = 0
             if isinstance(self.fusion_method, str):
                 fusion_methods = [self.fusion_method]
             else:
-                fusion_methods = self.fusion_method
+                fusion_methods = self.fusion_method  # type: ignore
             for fm in fusion_methods:
                 if fm == "concatenate":
                     output_dim += sum([model.output_dim for model in self.models])
@@ -291,11 +303,20 @@ def check_input_parameters(self):  # noqa: C901
         if isinstance(self.fusion_method, str):
             if not any(
                 x == self.fusion_method
-                for x in ["concatenate", "min", "max", "mean", "sum", "mult", "head"]
+                for x in [
+                    "concatenate",
+                    "min",
+                    "max",
+                    "mean",
+                    "sum",
+                    "dot",
+                    "mult",
+                    "head",
+                ]
             ):
                 raise ValueError(
-                    "fusion_method must be one of ['concatenate', 'mean', 'max', 'sum', 'mult', 'head'] "
-                    "or a list of those"
+                    "fusion_method must be one of ['concatenate', 'mean', 'max', 'sum', 'mult', 'dot', 'head'] "
+                    "or a list of any those but 'dot'"
                 )
 
             if (
@@ -323,14 +344,15 @@ def check_input_parameters(self):  # noqa: C901
                         "mean",
                         "sum",
                         "mult",
+                        "dot",
                         "head",
                     ]
                 )
                 for fm in self.fusion_method
             ):
                 raise ValueError(
-                    "fusion_method must be one of ['concatenate', 'mean', 'max', 'sum', 'mult', 'head'] "
-                    "or a list of those"
+                    "fusion_method must be one of ['concatenate', 'mean', 'max', 'sum', 'mult', 'dot', 'head'] "
+                    "or a list of those but 'dot'"
                 )
 
             if (
@@ -351,9 +373,11 @@ def check_input_parameters(self):  # noqa: C901
                         "projection_method must be one of ['min', 'max', 'mean']"
                     )
 
-        if "head" in self.fusion_method and isinstance(self.fusion_method, list):
+        if any(x in self.fusion_method for x in ["head", "dot"]) and isinstance(
+            self.fusion_method, list
+        ):
             raise ValueError(
-                "When using 'head' as fusion_method, no other method should be provided"
+                "When using 'head' or 'dot' as fusion_method, no other method should be provided"
             )
 
     def __repr__(self):

pytorch_widedeep/models/text/rnns/basic_rnn.py

@@ -2,7 +2,7 @@
 
 import numpy as np
 import torch
-from torch import nn, lstm
+from torch import nn
 
 from pytorch_widedeep.wdtypes import (
     List,

pytorch_widedeep/models/wide_deep.py

@@ -413,13 +413,13 @@ def _set_model_component(
         if isinstance(component, list):
             component_: Optional[Union[nn.ModuleList, WDModel]] = nn.ModuleList()
             for cp in component:
-                if self.with_deephead:
+                if self.with_deephead or cp.output_dim == 1:
                     component_.append(cp)
                 else:
                     component_.append(
                         nn.Sequential(cp, nn.Linear(cp.output_dim, self.pred_dim))
                     )
-        elif self.with_deephead:
+        elif self.with_deephead or component.output_dim == 1:
             component_ = component
         elif is_deeptabular and self.is_tabnet:
             component_ = nn.Sequential(
@@ -463,6 +463,11 @@ def _check_inputs(  # noqa: C901
             else:
                 if not hasattr(deeptabular, "output_dim"):
                     raise AttributeError(err_msg)
+                # the following assertion is thought for those cases where we
+                # use fusion with 'dot product' so that the output_dim will
+                # be 1 and the pred_dim is not 1
+                if deeptabular.output_dim == 1:
+                    assert pred_dim == 1, "If 'output_dim' is 1, 'pred_dim' must be 1"
 
         if deeptabular is not None:
             is_tabnet = False
@@ -502,6 +507,8 @@ def _check_inputs(  # noqa: C901
             else:
                 if not hasattr(deeptext, "output_dim"):
                     raise AttributeError(err_msg)
+                if deeptext.output_dim == 1:
+                    assert pred_dim == 1, "If 'output_dim' is 1, 'pred_dim' must be 1"
 
         if deepimage is not None:
             err_msg = "deepimage model must have an 'output_dim' attribute or property."
@@ -512,6 +519,8 @@ def _check_inputs(  # noqa: C901
             else:
                 if not hasattr(deepimage, "output_dim"):
                     raise AttributeError(err_msg)
+                if deepimage.output_dim == 1:
+                    assert pred_dim == 1, "If 'output_dim' is 1, 'pred_dim' must be 1"
 
         if deephead is not None and head_hidden_dims is not None:
             raise ValueError(

pytorch_widedeep/training/_trainer_utils.py

@@ -22,8 +22,10 @@
     FocalR_RMSELoss,
 )
 from pytorch_widedeep.wdtypes import (
+    Any,
     Dict,
     List,
+    Tuple,
     Union,
     Compose,
     Literal,
@@ -115,7 +117,7 @@ def wd_train_val_split(  # noqa: C901
     seed: int,
     method: Literal["regression", "binary", "multiclass", "qregression"],
     X_wide: Optional[np.ndarray] = None,
-    X_tab: Optional[np.ndarray] = None,
+    X_tab: Optional[Union[np.ndarray, List[np.ndarray]]] = None,
     X_text: Optional[Union[np.ndarray, List[np.ndarray]]] = None,
     X_img: Optional[Union[np.ndarray, List[np.ndarray]]] = None,
     X_train: Optional[Dict[str, Union[np.ndarray, List[np.ndarray]]]] = None,
@@ -174,6 +176,7 @@ def wd_train_val_split(  # noqa: C901
                 target is not None
             ), "if the validation split is specified, the target must also be specified"
             X_train = _build_train_dict(X_wide, X_tab, X_text, X_img, target)
+
         y_tr, y_val, idx_tr, idx_val = train_test_split(
             X_train["target"],
             np.arange(len(X_train["target"])),
@@ -187,46 +190,23 @@ def wd_train_val_split(  # noqa: C901
         )
         X_tr, X_val = {"target": y_tr}, {"target": y_val}
         if "X_wide" in X_train.keys():
-            X_tr["X_wide"], X_val["X_wide"] = (
-                X_train["X_wide"][idx_tr],
-                X_train["X_wide"][idx_val],
+            # the wide component will never be a list, but can still be passed
+            # to '_wd_train_val_split_component'
+            X_tr, X_val = _wd_train_val_split_component(
+                X_train, X_tr, X_val, idx_tr, idx_val, "X_wide"
             )
         if "X_tab" in X_train.keys():
-            X_tr["X_tab"], X_val["X_tab"] = (
-                X_train["X_tab"][idx_tr],
-                X_train["X_tab"][idx_val],
+            X_tr, X_val = _wd_train_val_split_component(
+                X_train, X_tr, X_val, idx_tr, idx_val, "X_tab"
             )
         if "X_text" in X_train.keys():
-            if isinstance(X_train["X_text"], list):
-                X_tr["X_text"], X_val["X_text"] = (
-                    [
-                        X_train["X_text"][i][idx_tr]
-                        for i in range(len(X_train["X_text"]))
-                    ],
-                    [
-                        X_train["X_text"][i][idx_val]
-                        for i in range(len(X_train["X_text"]))
-                    ],
-                )
-            else:
-                X_tr["X_text"], X_val["X_text"] = (
-                    X_train["X_text"][idx_tr],
-                    X_train["X_text"][idx_val],
-                )
+            X_tr, X_val = _wd_train_val_split_component(
+                X_train, X_tr, X_val, idx_tr, idx_val, "X_text"
+            )
         if "X_img" in X_train.keys():
-            if isinstance(X_train["X_img"], list):
-                X_tr["X_img"], X_val["X_img"] = (
-                    [X_train["X_img"][i][idx_tr] for i in range(len(X_train["X_img"]))],
-                    [
-                        X_train["X_img"][i][idx_val]
-                        for i in range(len(X_train["X_img"]))
-                    ],
-                )
-            else:
-                X_tr["X_img"], X_val["X_img"] = (
-                    X_train["X_img"][idx_tr],
-                    X_train["X_img"][idx_val],
-                )
+            X_tr, X_val = _wd_train_val_split_component(
+                X_train, X_tr, X_val, idx_tr, idx_val, "X_img"
+            )
         train_set = WideDeepDataset(**X_tr, transforms=transforms)  # type: ignore
         eval_set = WideDeepDataset(**X_val, transforms=transforms)  # type: ignore
     else:
@@ -239,9 +219,34 @@ def wd_train_val_split(  # noqa: C901
     return train_set, eval_set
 
 
+def _wd_train_val_split_component(
+    X: Dict[str, Union[np.ndarray, List[np.ndarray]]],
+    X_tr: Dict[str, Union[np.ndarray, List[np.ndarray]]],
+    X_val: Dict[str, Union[np.ndarray, List[np.ndarray]]],
+    idx_tr: Any,  # is a numpy array but sklearn's train_test_split returns a non-sensical type
+    idx_val: Any,
+    component_type: Literal["X_wide", "X_tab", "X_text", "X_img"],
+) -> Tuple[
+    Dict[str, Union[np.ndarray, List[np.ndarray]]],
+    Dict[str, Union[np.ndarray, List[np.ndarray]]],
+]:
+    if isinstance(X[component_type], list):
+        X_tr[component_type], X_val[component_type] = (
+            [X[component_type][i][idx_tr] for i in range(len(X[component_type]))],
+            [X[component_type][i][idx_val] for i in range(len(X[component_type]))],
+        )
+    else:
+        X_tr[component_type], X_val[component_type] = (
+            X[component_type][idx_tr],
+            X[component_type][idx_val],
+        )
+
+    return X_tr, X_val
+
+
 def _build_train_dict(
     X_wide: Optional[np.ndarray],
-    X_tab: Optional[np.ndarray],
+    X_tab: Optional[Union[np.ndarray, List[np.ndarray]]],
     X_text: Optional[Union[np.ndarray, List[np.ndarray]]],
     X_img: Optional[Union[np.ndarray, List[np.ndarray]]],
     target: np.ndarray,

pytorch_widedeep/training/_wd_dataset.py

@@ -32,7 +32,7 @@ class WideDeepDataset(Dataset):
     def __init__(
         self,
         X_wide: Optional[np.ndarray] = None,
-        X_tab: Optional[np.ndarray] = None,
+        X_tab: Optional[Union[np.ndarray, List[np.ndarray]]] = None,
         X_text: Optional[Union[np.ndarray, List[np.ndarray]]] = None,
         X_img: Optional[Union[np.ndarray, List[np.ndarray]]] = None,
         target: Optional[np.ndarray] = None,
@@ -60,7 +60,10 @@ def __getitem__(self, idx: int):  # noqa: C901
         if self.X_wide is not None:
             x.wide = self.X_wide[idx]
         if self.X_tab is not None:
-            x.deeptabular = self.X_tab[idx]
+            if isinstance(self.X_tab, list):
+                x.deeptabular = [self.X_tab[i][idx] for i in range(len(self.X_tab))]
+            else:
+                x.deeptabular = self.X_tab[idx]
         if self.X_text is not None:
             if isinstance(self.X_text, list):
                 x.deeptext = [self.X_text[i][idx] for i in range(len(self.X_text))]
@@ -112,7 +115,10 @@ def __len__(self):
         if self.X_wide is not None:
             return len(self.X_wide)
         if self.X_tab is not None:
-            return len(self.X_tab)
+            if isinstance(self.X_tab, list):
+                return len(self.X_tab[0])
+            else:
+                return len(self.X_tab)
         if self.X_text is not None:
             if isinstance(self.X_text, list):
                 return len(self.X_text[0])

pytorch_widedeep/training/trainer.py

@@ -270,7 +270,7 @@ def __init__(
     def fit(  # noqa: C901
         self,
         X_wide: Optional[np.ndarray] = None,
-        X_tab: Optional[np.ndarray] = None,
+        X_tab: Optional[Union[np.ndarray, List[np.ndarray]]] = None,
         X_text: Optional[Union[np.ndarray, List[np.ndarray]]] = None,
         X_img: Optional[Union[np.ndarray, List[np.ndarray]]] = None,
         X_train: Optional[Dict[str, Union[np.ndarray, List[np.ndarray]]]] = None,
@@ -529,7 +529,7 @@ def fit(  # noqa: C901
     def predict(  # type: ignore[override, return]
         self,
         X_wide: Optional[np.ndarray] = None,
-        X_tab: Optional[np.ndarray] = None,
+        X_tab: Optional[Union[np.ndarray, List[np.ndarray]]] = None,
         X_text: Optional[Union[np.ndarray, List[np.ndarray]]] = None,
         X_img: Optional[Union[np.ndarray, List[np.ndarray]]] = None,
         X_test: Optional[Dict[str, Union[np.ndarray, List[np.ndarray]]]] = None,
@@ -585,7 +585,7 @@ def predict(  # type: ignore[override, return]
     def predict_uncertainty(  # type: ignore[return]
         self,
         X_wide: Optional[np.ndarray] = None,
-        X_tab: Optional[np.ndarray] = None,
+        X_tab: Optional[Union[np.ndarray, List[np.ndarray]]] = None,
         X_text: Optional[Union[np.ndarray, List[np.ndarray]]] = None,
         X_img: Optional[Union[np.ndarray, List[np.ndarray]]] = None,
         X_test: Optional[Dict[str, Union[np.ndarray, List[np.ndarray]]]] = None,
@@ -682,7 +682,7 @@ def predict_uncertainty(  # type: ignore[return]
     def predict_proba(  # type: ignore[override, return]  # noqa: C901
         self,
         X_wide: Optional[np.ndarray] = None,
-        X_tab: Optional[np.ndarray] = None,
+        X_tab: Optional[Union[np.ndarray, List[np.ndarray]]] = None,
         X_text: Optional[Union[np.ndarray, List[np.ndarray]]] = None,
         X_img: Optional[Union[np.ndarray, List[np.ndarray]]] = None,
         X_test: Optional[Dict[str, Union[np.ndarray, List[np.ndarray]]]] = None,
@@ -942,7 +942,7 @@ def _eval_step(
                     X[k] = v.to(self.device)
             y = (
                 target.view(-1, 1).float()
-                if self.method not in ["multiclass", "qregression"]
+                if self.method not in ["multiclass", "qregression", "multitarget"]
                 else target
             )
             y = y.to(self.device)
@@ -971,14 +971,15 @@ def _get_score(self, y_pred, y):
                 score = self.metric(y_pred, y)
             if self.method == "multiclass":
                 score = self.metric(F.softmax(y_pred, dim=1), y)
+            # TO DO: handle multitarget
             return score
         else:
             return None
 
     def _predict(  # type: ignore[override, return]  # noqa: C901
         self,
         X_wide: Optional[np.ndarray] = None,
-        X_tab: Optional[np.ndarray] = None,
+        X_tab: Optional[Union[np.ndarray, List[np.ndarray]]] = None,
         X_text: Optional[Union[np.ndarray, List[np.ndarray]]] = None,
         X_img: Optional[Union[np.ndarray, List[np.ndarray]]] = None,
         X_test: Optional[Dict[str, Union[np.ndarray, List[np.ndarray]]]] = None,