-- added save and load model unit test

Author: manujosephv

tests/test_common.py

@@ -0,0 +1,141 @@
+#!/usr/bin/env python
+"""Tests for `pytorch_tabular` package."""
+
+import pytest
+import numpy as np
+import torch
+from sklearn.preprocessing import PowerTransformer
+from pytorch_tabular.config import DataConfig, OptimizerConfig, TrainerConfig
+from pytorch_tabular.models import CategoryEmbeddingModelConfig, AutoIntConfig, NodeConfig, TabNetModelConfig, CategoryEmbeddingMDNConfig
+from pytorch_tabular import TabularModel
+from pytorch_tabular.categorical_encoders import CategoricalEmbeddingTransformer
+
+MODEL_CONFIGS = [
+    CategoryEmbeddingModelConfig, 
+    AutoIntConfig, 
+    NodeConfig, 
+    TabNetModelConfig, 
+    CategoryEmbeddingMDNConfig
+]
+
+
+def fake_metric(y_hat, y):
+    return (y_hat - y).mean()
+
+
+@pytest.mark.parametrize(
+    "model_config_class",
+    MODEL_CONFIGS,
+)
+@pytest.mark.parametrize(
+    "continuous_cols",
+    [
+        [
+            "AveRooms",
+            "AveBedrms",
+            "Population",
+            "AveOccup",
+            "Latitude",
+            "Longitude",
+        ],
+    ],
+)
+@pytest.mark.parametrize("categorical_cols", [["HouseAgeBin"]])
+@pytest.mark.parametrize("custom_metrics", [None, [fake_metric]])
+@pytest.mark.parametrize("custom_loss", [None, torch.nn.L1Loss()])
+@pytest.mark.parametrize("custom_optimizer", [None, torch.optim.Adagrad])
+def test_save_load(
+    regression_data,
+    model_config_class,
+    continuous_cols,
+    categorical_cols,
+    custom_metrics,
+    custom_loss,
+    custom_optimizer,
+    tmpdir
+):
+    (train, test, target) = regression_data
+    data_config = DataConfig(
+        target=target,
+        continuous_cols=continuous_cols,
+        categorical_cols=categorical_cols,
+    )
+    model_config_params = dict(task="regression")
+    model_config = model_config_class(**model_config_params)
+    trainer_config = TrainerConfig(
+        max_epochs=3, checkpoints=None, early_stopping=None, gpus=0
+    )
+    optimizer_config = OptimizerConfig()
+
+    tabular_model = TabularModel(
+        data_config=data_config,
+        model_config=model_config,
+        optimizer_config=optimizer_config,
+        trainer_config=trainer_config,
+    )
+    tabular_model.fit(
+        train=train,
+        test=test,
+        metrics=custom_metrics,
+        loss=custom_loss,
+        optimizer=custom_optimizer,
+        optimizer_params=None if custom_optimizer is None else {},
+    )
+
+    result_1 = tabular_model.evaluate(test)
+    print(result_1)
+    tmpdir.mkdir("save_model")
+    tabular_model.save_model("save_model")
+    new_mdl = TabularModel.load_from_checkpoint("save_model")
+    result_2 = new_mdl.evaluate(test)
+    assert result_1[0][f'test_{tabular_model.model.hparams.metrics[0]}'] == result_2[0][f'test_{new_mdl.model.hparams.metrics[0]}']
+
+# import numpy as np
+# import pandas as pd
+# from sklearn.datasets import fetch_california_housing, fetch_covtype
+# from pathlib import Path
+
+# def regression_data():
+#     dataset = fetch_california_housing(data_home="data", as_frame=True)
+#     df = dataset.frame.sample(5000)
+#     df["HouseAgeBin"] = pd.qcut(df["HouseAge"], q=4)
+#     df["HouseAgeBin"] = "age_" + df.HouseAgeBin.cat.codes.astype(str)
+#     test_idx = df.sample(int(0.2 * len(df)), random_state=42).index
+#     test = df[df.index.isin(test_idx)]
+#     train = df[~df.index.isin(test_idx)]
+#     return (train, test, dataset.target_names)
+
+
+# def classification_data():
+#     dataset = fetch_covtype(data_home="data")
+#     data = np.hstack([dataset.data, dataset.target.reshape(-1, 1)])[:10000, :]
+#     col_names = [f"feature_{i}" for i in range(data.shape[-1])]
+#     col_names[-1] = "target"
+#     data = pd.DataFrame(data, columns=col_names)
+#     data["feature_0_cat"] = pd.qcut(data["feature_0"], q=4)
+#     data["feature_0_cat"] = "feature_0_" + data.feature_0_cat.cat.codes.astype(str)
+#     test_idx = data.sample(int(0.2 * len(data)), random_state=42).index
+#     test = data[data.index.isin(test_idx)]
+#     train = data[~data.index.isin(test_idx)]
+#     return (train, test, ["target"])
+
+
+# test_save_load(
+#     regression_data(),
+#     model_config_class=CategoryEmbeddingModelConfig,
+#     continuous_cols=[
+#         "AveRooms",
+#         "AveBedrms",
+#         "Population",
+#         "AveOccup",
+#         "Latitude",
+#         "Longitude",
+#     ],
+#     categorical_cols=[],
+#     custom_metrics = None, #[fake_metric],
+#     custom_loss = None, custom_optimizer = None,
+#     tmpdir = Path("tmp")
+# )
+# test_embedding_transformer(regression_data())
+
+# classification_data()