feat(skorch): add an inherited class from skorch.NeuralNet that is compatible with PyTorch Frame

CHANGELOG.md

@@ -22,6 +22,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Added `MovieLens 1M` dataset ([#397](https://github.com/pyg-team/pytorch-frame/pull/397))
 - Added light-weight MLP ([#372](https://github.com/pyg-team/pytorch-frame/pull/372))
+- Added an inherited class from skorch.NeuralNet that is compatible with PyTorch Frame ([#375](https://github.com/pyg-team/pytorch-frame/pull/375))
 - Added R^2 metric ([#403](https://github.com/pyg-team/pytorch-frame/pull/403))
 
 ### Changed

README.md

@@ -68,12 +68,15 @@ PyTorch Frame democratizes deep learning research for tabular data, catering to
 
 <hr style="border: 0.5px solid #ccc;">
 
-* [Library Highlights](#library-highlights)
-* [Architecture Overview](#architecture-overview)
-* [Quick Tour](#quick-tour)
-* [Implemented Deep Tabular Models](#implemented-deep-tabular-models)
-* [Benchmark](#benchmark)
-* [Installation](#installation)
+- [Library Highlights](#library-highlights)
+- [Architecture Overview](#architecture-overview)
+- [Quick Tour](#quick-tour)
+  - [Build and train your own deep tabular model](#build-and-train-your-own-deep-tabular-model)
+  - [Scikit-learn Compatible API (Experimental)](#scikit-learn-compatible-api-experimental)
+- [Implemented Deep Tabular Models](#implemented-deep-tabular-models)
+- [Benchmark](#benchmark)
+- [Installation](#installation)
+- [Cite](#cite)
 
 ## Library Highlights
 
@@ -206,6 +209,67 @@ for epoch in range(50):
         loss.backward()
 ```
 
+### Scikit-learn Compatible API (Experimental)
+
+A scikit-learn compliant API based on skorch allows DataFrame to be trained directly as a dataset. However, this has many limitations.
+
+```python
+from typing import Any
+
+import torch.nn as nn
+from sklearn.datasets import load_diabetes
+from sklearn.metrics import mean_squared_error
+from sklearn.model_selection import train_test_split
+
+from torch_frame import stype
+from torch_frame.data.stats import StatType
+from torch_frame.nn import Trompt
+from torch_frame.nn.models.trompt import Trompt
+from torch_frame.utils.skorch import NeuralNetPytorchFrame
+
+# load the diabetes dataset
+X, y = load_diabetes(return_X_y=True, as_frame=True)
+
+# split the data into training and testing sets
+X_train, X_test, y_train, y_test = train_test_split(X, y)
+
+
+# define the function to get the module
+def get_module(*, col_stats: dict[str, dict[StatType, Any]],
+               col_names_dict: dict[stype, list[str]]) -> Trompt:
+    channels = 8
+    out_channels = 1
+    num_prompts = 2
+    num_layers = 3
+    return Trompt(channels=channels, out_channels=out_channels,
+                  num_prompts=num_prompts, num_layers=num_layers,
+                  col_stats=col_stats, col_names_dict=col_names_dict,
+                  stype_encoder_dicts=None)
+
+
+# wrap the function in a NeuralNetPytorchFrame
+# NeuralNetClassifierPytorchFrame and NeuralNetBinaryClassifierPytorchFrame
+# are also available
+net = NeuralNetPytorchFrame(
+    module=get_module,
+    criterion=nn.MSELoss(),
+    max_epochs=10,
+    verbose=1,
+    lr=0.0001,
+    batch_size=30,
+)
+
+# fit the model
+net.fit(X_train, y_train)
+
+# predict on the test set
+y_pred = net.predict(X_test)
+
+# calculate the mean squared error
+mse = mean_squared_error(y_test, y_pred)
+print(mse)
+```
+
 ## Implemented Deep Tabular Models
 
 We list currently supported deep tabular models:

examples/revisiting.py

@@ -17,7 +17,10 @@
 import argparse
 import os.path as osp
 
+import numpy as np
+import pandas as pd
 import torch
+import torch.nn as nn
 import torch.nn.functional as F
 from tqdm import tqdm
 
@@ -46,6 +49,8 @@
 parser.add_argument('--epochs', type=int, default=100)
 parser.add_argument('--seed', type=int, default=0)
 parser.add_argument('--compile', action='store_true')
+parser.add_argument("--framework", type=str, default="torch",
+                    choices=["torch", "skorch-dataframe"])
 args = parser.parse_args()
 
 torch.manual_seed(args.seed)
@@ -156,30 +161,82 @@ def test(loader: DataLoader) -> float:
         return rmse
 
 
-if is_classification:
-    metric = 'Acc'
-    best_val_metric = 0
-    best_test_metric = 0
-else:
-    metric = 'RMSE'
-    best_val_metric = float('inf')
-    best_test_metric = float('inf')
-
-for epoch in range(1, args.epochs + 1):
-    train_loss = train(epoch)
-    train_metric = test(train_loader)
-    val_metric = test(val_loader)
-    test_metric = test(test_loader)
-
-    if is_classification and val_metric > best_val_metric:
-        best_val_metric = val_metric
-        best_test_metric = test_metric
-    elif not is_classification and val_metric < best_val_metric:
-        best_val_metric = val_metric
-        best_test_metric = test_metric
-
-    print(f'Train Loss: {train_loss:.4f}, Train {metric}: {train_metric:.4f}, '
-          f'Val {metric}: {val_metric:.4f}, Test {metric}: {test_metric:.4f}')
-
-print(f'Best Val {metric}: {best_val_metric:.4f}, '
-      f'Best Test {metric}: {best_test_metric:.4f}')
+if args.framework == "torch":
+    if is_classification:
+        metric = 'Acc'
+        best_val_metric = 0
+        best_test_metric = 0
+    else:
+        metric = 'RMSE'
+        best_val_metric = float('inf')
+        best_test_metric = float('inf')
+
+    for epoch in range(1, args.epochs + 1):
+        train_loss = train(epoch)
+        train_metric = test(train_loader)
+        val_metric = test(val_loader)
+        test_metric = test(test_loader)
+
+        if is_classification and val_metric > best_val_metric:
+            best_val_metric = val_metric
+            best_test_metric = test_metric
+        elif not is_classification and val_metric < best_val_metric:
+            best_val_metric = val_metric
+            best_test_metric = test_metric
+
+        print(
+            f'Train Loss: {train_loss:.4f}, '
+            f'Train {metric}: {train_metric:.4f}, '
+            f'Val {metric}: {val_metric:.4f}, Test {metric}: {test_metric:.4f}'
+        )
+
+    print(f'Best Val {metric}: {best_val_metric:.4f}, '
+          f'Best Test {metric}: {best_test_metric:.4f}')
+elif args.framework == "skorch-dataframe":
+
+    from torch_frame.utils.skorch import (
+        NeuralNetClassifierPytorchFrame,
+        NeuralNetPytorchFrame,
+    )
+
+    df = dataset.df
+    df_train = pd.concat([train_dataset.df, val_dataset.df])
+    X_train, y_train = df_train.drop(
+        columns=[dataset.target_col, dataset.split_col]), df_train[
+            dataset.target_col]
+    df_test = test_dataset.df
+    X_test, y_test = df_test.drop(
+        columns=[dataset.target_col, dataset.split_col]), df_test[
+            dataset.target_col]
+
+    # use DataFrames with no `split_col` or `target_col`
+    # like normal sklearn datasets from now on
+    if is_classification:
+        net = NeuralNetClassifierPytorchFrame(
+            module=model,
+            criterion=nn.CrossEntropyLoss,
+            max_epochs=args.epochs,
+            lr=args.lr,
+            device=device,
+            verbose=1,
+            batch_size=args.batch_size,
+        )
+    else:
+        net = NeuralNetPytorchFrame(
+            module=model,
+            criterion=nn.MSELoss,
+            max_epochs=args.epochs,
+            lr=args.lr,
+            device=device,
+            verbose=1,
+            batch_size=args.batch_size,
+        )
+    net.fit(X_train, y_train)
+    y_pred = net.predict(X_test)
+
+    if is_classification:
+        test_acc = (y_pred.argmax(-1) == y_test).mean()
+        print(f"Test Acc: {test_acc:.4f}")
+    else:
+        test_rmse = np.sqrt(((y_pred.squeeze() - y_test)**2).mean())
+        print(f"Test RMSE: {test_rmse:.4f}")

examples/tutorial.py

@@ -7,7 +7,9 @@
 import os.path as osp
 from typing import Any, Dict, List
 
+import pandas as pd
 import torch
+import torch.nn as nn
 import torch.nn.functional as F
 from torch import Tensor
 from torch.nn import LayerNorm, Linear, Module, ModuleList
@@ -34,6 +36,8 @@
 parser.add_argument('--lr', type=float, default=0.001)
 parser.add_argument('--epochs', type=int, default=10)
 parser.add_argument('--seed', type=int, default=0)
+parser.add_argument("--framework", type=str, default="torch",
+                    choices=["torch", "skorch", "skorch-dataframe"])
 args = parser.parse_args()
 
 torch.manual_seed(args.seed)
@@ -223,7 +227,7 @@ def train(epoch: int) -> float:
     model.train()
     loss_accum = total_count = 0
 
-    for tf in tqdm(train_loader, desc=f'Epoch: {epoch}'):
+    for tf in tqdm(train_loader, desc=f"Epoch: {epoch}"):
         tf = tf.to(device)
         pred = model(tf)
         loss = F.cross_entropy(pred, tf.y)
@@ -250,17 +254,68 @@ def test(loader: DataLoader) -> float:
     return accum / total_count
 
 
-best_val_acc = 0
-best_test_acc = 0
-for epoch in range(1, args.epochs + 1):
-    train_loss = train(epoch)
-    train_acc = test(train_loader)
-    val_acc = test(val_loader)
-    test_acc = test(test_loader)
-    if best_val_acc < val_acc:
-        best_val_acc = val_acc
-        best_test_acc = test_acc
-    print(f'Train Loss: {train_loss:.4f}, Train Acc: {train_acc:.4f}, '
-          f'Val Acc: {val_acc:.4f}, Test Acc: {test_acc:.4f}')
-
-print(f'Best Val Acc: {best_val_acc:.4f}, Best Test Acc: {best_test_acc:.4f}')
+if args.framework == "torch":
+    best_val_acc = 0
+    best_test_acc = 0
+    for epoch in range(1, args.epochs + 1):
+        train_loss = train(epoch)
+        train_acc = test(train_loader)
+        val_acc = test(val_loader)
+        test_acc = test(test_loader)
+        if best_val_acc < val_acc:
+            best_val_acc = val_acc
+            best_test_acc = test_acc
+        print(f"Train Loss: {train_loss:.4f}, Train Acc: {train_acc:.4f}, "
+              f"Val Acc: {val_acc:.4f}, Test Acc: {test_acc:.4f}")
+
+    print(
+        f"Best Val Acc: {best_val_acc:.4f}, Best Test Acc: {best_test_acc:.4f}"
+    )
+elif args.framework == "skorch":
+
+    from torch_frame.utils.skorch import NeuralNetClassifierPytorchFrame
+
+    net = NeuralNetClassifierPytorchFrame(
+        module=model,
+        criterion=nn.CrossEntropyLoss,
+        max_epochs=args.epochs,
+        lr=args.lr,
+        device=device,
+        verbose=1,
+        batch_size=args.batch_size,
+    )
+    net.fit(dataset)
+    y_pred = net.predict(test_dataset)
+    test_acc = (torch.Tensor(y_pred).argmax(
+        dim=-1) == test_tensor_frame.y).float().mean()
+    print(f"Test Acc: {test_acc:.4f}")
+elif args.framework == "skorch-dataframe":
+
+    from torch_frame.utils.skorch import NeuralNetClassifierPytorchFrame
+
+    df = dataset.df
+    df_train = pd.concat([train_dataset.df, val_dataset.df])
+    X_train, y_train = df_train.drop(
+        columns=[dataset.target_col, dataset.split_col]), df_train[
+            dataset.target_col]
+    df_test = test_dataset.df
+    X_test, y_test = df_test.drop(
+        columns=[dataset.target_col, dataset.split_col]), df_test[
+            dataset.target_col]
+
+    # use DataFrames with no `split_col` or `target_col`
+    # like normal sklearn datasets from now on
+    net = NeuralNetClassifierPytorchFrame(
+        module=model,
+        criterion=nn.CrossEntropyLoss,
+        max_epochs=args.epochs,
+        lr=args.lr,
+        device=device,
+        verbose=1,
+        col_to_stype={"C_feature_7": stype.categorical},
+        batch_size=args.batch_size,
+    )
+    net.fit(X_train, y_train)
+    y_pred = net.predict(X_test)
+    test_acc = (y_pred.argmax(-1) == y_test).mean()
+    print(f"Test Acc: {test_acc:.4f}")

pyproject.toml

@@ -56,6 +56,7 @@ full=[
     "lightgbm",
     "datasets",
     "torchmetrics",
+    "skorch",
 ]
 
 [project.urls]