TabNet Classifier

TabNet model class and training logic.

class pytorch_tabnet.tab_model.TabNetClassifier(n_d: int = 8, n_a: int = 8, n_steps: int = 3, gamma: float = 1.3, cat_idxs: ~typing.List[int] = <factory>, cat_dims: ~typing.List[int] = <factory>, cat_emb_dim: int | ~typing.List[int] = 1, n_independent: int = 2, n_shared: int = 2, epsilon: float = 1e-15, momentum: float = 0.02, lambda_sparse: float = 0.001, seed: int = 0, clip_value: int = 1, verbose: int = 1, optimizer_fn: ~typing.Any = <class 'torch.optim.adam.Adam'>, optimizer_params: ~typing.Dict = <factory>, scheduler_fn: ~typing.Any = None, scheduler_params: ~typing.Dict = <factory>, mask_type: str = 'sparsemax', input_dim: int = None, output_dim: int = None, device_name: str = 'auto', n_shared_decoder: int = 1, n_indep_decoder: int = 1, grouped_features: ~typing.List[~typing.List[int]] = <factory>, compile_backend: str = '')[source]

Bases: TabSupervisedModel

TabNet model for classification tasks.

compute_loss(y_pred: Tensor, y_true: Tensor, w: Tensor | None = None) → Tensor[source]

Compute the loss for classification.

Parameters:

  • y_pred (torch.Tensor) – Network output.

  • y_true (torch.Tensor) – True labels.

  • w (Optional[torch.Tensor]) – Optional sample weights.

Returns:

Loss value.

Return type:

torch.Tensor

output_dim: int = None
predict_func(outputs: ndarray) → ndarray[source]

Convert network outputs to class predictions.

Parameters:

outputs (np.ndarray) – Network outputs.

Returns:

Predicted classes.

Return type:

np.ndarray

predict_proba(X: Tensor | ndarray) → ndarray[source]

Predict class probabilities for classification.

Parameters:

X (torch.Tensor or np.ndarray) – Input data.

Returns:

Probability predictions.

Return type:

np.ndarray

prepare_target(y: ndarray) → ndarray[source]

Map targets using the target mapper.

Parameters:

y (np.ndarray) – Target array.

Returns:

Mapped target array.

Return type:

np.ndarray

set_fit_request(*, X_train: bool | None | str = '$UNCHANGED$', batch_size: bool | None | str = '$UNCHANGED$', callbacks: bool | None | str = '$UNCHANGED$', compute_importance: bool | None | str = '$UNCHANGED$', drop_last: bool | None | str = '$UNCHANGED$', eval_metric: bool | None | str = '$UNCHANGED$', eval_name: bool | None | str = '$UNCHANGED$', eval_set: bool | None | str = '$UNCHANGED$', from_unsupervised: bool | None | str = '$UNCHANGED$', loss_fn: bool | None | str = '$UNCHANGED$', max_epochs: bool | None | str = '$UNCHANGED$', num_workers: bool | None | str = '$UNCHANGED$', patience: bool | None | str = '$UNCHANGED$', pin_memory: bool | None | str = '$UNCHANGED$', virtual_batch_size: bool | None | str = '$UNCHANGED$', warm_start: bool | None | str = '$UNCHANGED$', weights: bool | None | str = '$UNCHANGED$', y_train: bool | None | str = '$UNCHANGED$') → TabNetClassifier

Request metadata passed to the fit method.

Note that this method is only relevant if enable_metadata_routing=True (see sklearn.set_config()). Please see User Guide on how the routing mechanism works.

The options for each parameter are:

  • True: metadata is requested, and passed to fit if provided. The request is ignored if metadata is not provided.

  • False: metadata is not requested and the meta-estimator will not pass it to fit.

  • None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.

  • str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Note

This method is only relevant if this estimator is used as a sub-estimator of a meta-estimator, e.g. used inside a Pipeline. Otherwise it has no effect.

Parameters:

  • X_train (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for X_train parameter in fit.

  • batch_size (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for batch_size parameter in fit.

  • callbacks (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for callbacks parameter in fit.

  • compute_importance (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for compute_importance parameter in fit.

  • drop_last (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for drop_last parameter in fit.

  • eval_metric (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for eval_metric parameter in fit.

  • eval_name (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for eval_name parameter in fit.

  • eval_set (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for eval_set parameter in fit.

  • from_unsupervised (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for from_unsupervised parameter in fit.

  • loss_fn (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for loss_fn parameter in fit.

  • max_epochs (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for max_epochs parameter in fit.

  • num_workers (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for num_workers parameter in fit.

  • patience (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for patience parameter in fit.

  • pin_memory (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for pin_memory parameter in fit.

  • virtual_batch_size (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for virtual_batch_size parameter in fit.

  • warm_start (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for warm_start parameter in fit.

  • weights (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for weights parameter in fit.

  • y_train (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for y_train parameter in fit.

Returns:

self – The updated object.

Return type:

object

stack_batches(list_y_true: List[Tensor], list_y_score: List[Tensor]) → Tuple[Tensor, Tensor][source]

Stack batches of true and predicted values.

Parameters:

  • list_y_true (List[torch.Tensor]) – List of true labels for each batch.

  • list_y_score (List[torch.Tensor]) – List of predicted scores for each batch.

Returns:

Stacked true labels and predicted scores.

Return type:

Tuple[torch.Tensor, torch.Tensor]

update_fit_params(X_train: ndarray, y_train: ndarray, eval_set: List[Tuple[ndarray, ndarray]], weights: bool | Dict[str, Any]) → None[source]

Update fit parameters for classification.

Parameters:

  • X_train (np.ndarray) – Training data.

  • y_train (np.ndarray) – Training targets.

  • eval_set (list) – List of evaluation sets.

  • weights (bool or dict) – Class weights.

weight: Any = 0
weight_updater(weights: bool | Dict[str | int, Any] | Any) → bool | Dict[str | int, Any][source]

Update class weights for training.

Parameters:

weights (bool, dict, or any) – Class weights or indicator.

Returns:

Updated weights.

Return type:

bool or dict

Example

from pytorch_tabnet.tab_model import TabNetClassifier
import numpy as np
X = np.random.rand(100, 10)
y = np.random.randint(0, 2, size=(100,))
clf = TabNetClassifier()
clf.fit(X_train=X, y_train=y)
preds = clf.predict(X)