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

CHANGELOG.md

@@ -9,6 +9,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 ### Added
 
 - 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))
 
 ### Changed
 

examples/revisiting.py

@@ -46,6 +46,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 +158,87 @@ 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":
+    import numpy as np
+    import pandas as pd
+    import torch.nn as nn
+
+    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,
+            # col_to_stype={"C_feature_7": stype.categorical},
+            batch_size=args.batch_size,
+        )
+    else:
+        net = NeuralNetPytorchFrame(
+            module=model,
+            criterion=nn.MSELoss,
+            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)
+
+    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

@@ -34,6 +34,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 +225,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 +252,71 @@ 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":
+    import torch.nn as nn
+
+    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":
+    import pandas as pd
+    import torch.nn as nn
+
+    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}")