repository·main·Indexed 18 days ago
https://github.com/dholzmueller/pytabkitA library providing scikit-learn interfaces for modern tabular machine learning models (classification and regression) and benchmarking tools. It features highly tuned models like RealMLP, TabM, and TabR, and supports advanced techniques including hyperparameter optimization (HPO), post-hoc calibration, and weighted ensembling. The library includes the pytabkit.models module for ML functionality and the pytabkit.bench module for systematic benchmarking of tabular models.
Due to the vectorization of Neural Network models, PyTabKit uses specific terms for split counts. Understanding these is crucial when configuring models:
n_cv: Number of training-validation splits in cross-validation (bagging).n_refit: Number of models refitted on training+validation data after the CV stage.n_tv_splits (or n_models): Number of training-validation splits used in the current training (can be n_cv or n_refit).n_tt_splits (or n_parallel): Number of trainval-test splits used. This is typically 1 when using the scikit-learn interface, but can be larger when using RealMLP through the benchmark.The pytabkit.bench module is designed for systematic benchmarking of tabular models. Key capabilities include:
PyTabKit handles hyperparameters in two ways:
**config (or **kwargs).Warning: Because parameters are passed through nested functions via **config, typos in parameter names will not be caught by the interpreter and may lead to parameters being ignored or passed incorrectly.
To avoid confusion, PyTabKit uses unique names for similar parameters. For example:
opt_eps: Epsilon parameter for the optimizer.ls_eps: Epsilon parameter for label smoothing.PyTabKit introduces a Variable class (inheriting from torch.nn.Parameter) to manage parameters with metadata:
Variable has a trainable: bool parameter. If False, it is registered using register_buffer().'net/first_layer/layer-0/weight'). This is useful for regex-based hyperparameter assignment.**kwargs of NNAlgInterface.Example of assigning a specific learning rate to a layer via regex:
# Example pattern for NNAlgInterface kwargs
lr={'': global_lr, '.*first_layer.*weight': first_layer_weight_lr}For advanced usage, especially for benchmarking, use the AlgInterface (found in alg_interfaces/alg_interfaces.py) instead of the scikit-learn wrappers. AlgInterface provides features that the scikit-learn interfaces lack:
RealMLP-TD and RealMLP-TD-S).n_threads, and time.The pytabkit.models module contains the core machine learning functionality. You can use it to:
Datasets are represented internally using the DictDataset class, which contains a dictionary of PyTorch tensors. The standard keys are:
'x_cont': Continuous features.'x_cat': Categorical features (with dtype=torch.long).'y': Labels.Each key in the dictionary is accompanied by a TensorInfo object in a tensor_infos dictionary. This object describes the number of features and, for categorical variables or classification labels, the number of categories.
Important Notes:
0 is reserved for missing values or values not seen during training.PyTabKit uses a specialized class structure designed for vectorized Neural Networks. Unlike standard PyTorch nn.Module patterns where preprocessing is separate, PyTabKit unifies preprocessing and NN layers into a single structure. This allows multiple NNs to share a single non-preprocessed dataset in GPU RAM while maintaining different preprocessing parameters (fitted on different data splits), significantly saving GPU memory.
The architecture relies on three core base classes defined in model/base.py:
Layer: Similar to nn.Module but does not perform random initialization in the constructor. Instead, it accepts pre-initialized parameters. Layers work with DictDataset (containing x_cont and x_cat tensors) and can process labels y to implement features like mixup or label smoothing.Layer.stack().Fitter: Responsible for initializing the NN based on a single forward pass on training/validation sets. It uses fit() or fit_transform() (similar to scikit-learn) to return a Layer object.FitterFactory: An architecture builder that creates Fitter objects based on the input/output shapes and types (provided via tensor_infos). It handles logic such as choosing between one-hot encoding or Embedding layers based on category sizes.TabNN models in PyTabKit are implemented using PyTorch Lightning. To train them directly, you must follow a specific workflow involving DictDataset for data handling, TabNNModule for the model definition, and InterfaceResources for resource management.
Key implementation requirements:
DictDataset to wrap your tensors. The models have specific requirements for dataloaders that necessitate this format.InterfaceResources object rather than letting Lightning manage them automatically.Trainer, certain parameters are obligatory for TabNNModule to function correctly (e.g., specific callbacks and logger settings).from sklearn.datasets import make_classification
from sklearn.model_selection import train_test_split
from pytabkit.models.alg_interfaces.base import SplitIdxs, InterfaceResources
from pytabkit.models.data.data import DictDataset, TensorInfo
from pytabkit.models.sklearn.default_params import DefaultParams
from pytabkit.models.training.lightning_modules import TabNNModule
import lightning.pytorch as pl
import numpy as np
import torch
# 1. Setup Data
X, y = make_classification()
idxs = np.arange(len(X))
trainval_idxs, test_idxs = train_test_split(idxs, test_size=0.2)
# Create DictDataset
ds = DictDataset(
tensors={
'x_cont': torch.as_tensor(X, dtype=torch.float32),
'x_cat': torch.zeros(len(X), 0),
'y': torch.as_tensor(y, dtype=torch.long)[:, None]
},
tensor_infos={
'x_cont': TensorInfo(feat_shape=[X.shape[1]]),
'x_cat': TensorInfo(cat_sizes=[]),
'y': TensorInfo(cat_sizes=[np.max(y) + 1])
}
)
# Define splitting indices
train_val_splitting_idxs_list = [
SplitIdxs(
train_idxs=torch.as_tensor(np.stack([trainval_idxs], axis=0), dtype=torch.long),
val_idxs=torch.as_tensor(np.stack([trainval_idxs], axis=0), dtype=torch.long),
test_idxs=torch.as_tensor(test_idxs, dtype=torch.long),
split_seed=0,
sub_split_seeds=[0],
split_id=0
)
]
test_ds = ds.get_sub_dataset(torch.as_tensor(test_idxs, dtype=torch.long))
# 2. Setup Resources
interface_resources = InterfaceResources(n_threads=4, gpu_devices=[])
# 3. Define and Compile Model
# Use DefaultParams for specific architectures like RealMLP_TD_CLASS
nn_model = TabNNModule(**DefaultParams.RealMLP_TD_CLASS)
nn_model.compile_model(ds, train_val_splitting_idxs_list, interface_resources)
# 4. Train with Lightning Trainer
trainer = pl.Trainer(
callbacks=nnn_model.create_callbacks(),
max_epochs=200,
enable_checkpointing=False,
enable_progress_bar=False,
num_sanity_val_steps=0,
logger=pl.loggers.logger.DummyLogger(),
)
trainer.fit(
model=nnn_model,
train_dataloaders=nnn_model.train_dl,
val_dataloaders=nnn_model.val_dl
)
# 5. Predict
pred = trainer.predict(
model=nnn_model,
dataloaders=nnn_model.get_predict_dataloader(test_ds)
)Most models in PyTabKit follow the scikit-learn API (fit, predict, predict_proba). The models automatically handle GPU selection, categorical column detection, and preprocessing (numerical variables and regression targets).
Note: Missing numerical values are currently not allowed and must be imputed before passing data to the model.
from pytabkit import RealMLP_TD_Classifier
model = RealMLP_TD_Classifier()
model.fit(X_train, y_train)
predictions = model.predict(X_test)RealMLP and several other models can be saved using standard Python pickling (or dill).
GPU to CPU Migration:
If a model was trained on a GPU, standard pickling will attempt to restore it to a GPU. To load a GPU-trained model onto a CPU, use torch.load with map_location='cpu'.
Warning: While this allows loading, calling .predict() on a model loaded this way may fail due to pytorch-lightning device issues.
import torch
import dill
# Saving a model
torch.save(model, 'model.pkl', pickle_module=dill, _use_new_zipfile_serialization=False)
# Loading a GPU-trained model to CPU
model = torch.load('model.pkl', map_location='cpu', pickle_module=dill)