说明文档

TabNet：专注且可解释的表格数据学习

这是TabNet的PyTorch实现（Arik, S. O., & Pfister, T. (2019). TabNet: Attentive Interpretable Tabular Learning. arXiv预印本arXiv:1908.07442）。https://arxiv.org/pdf/1908.07442.pdf。请注意，为了改进该库，我们在此过程中做出了一些不同的选择，这些选择可能与原始论文有所不同。

有任何问题吗？想贡献代码吗？想和我们交流吗？您可以通过Slack加入我们。

安装

简单安装

您可以使用pip或conda进行安装，如下所示。

使用pip

pip install pytorch-tabnet

使用conda

conda install -c conda-forge pytorch-tabnet

源代码

如果您想在本地Docker容器中使用它：

• git clone git@github.com:dreamquark-ai/tabnet.git

• cd tabnet进入仓库目录

---

仅CPU

• make start构建并进入容器

GPU

• make start-gpu构建并进入GPU容器

---

• poetry install安装所有依赖项，包括Jupyter

• 在同一终端中运行make notebook。然后您可以按照链接打开一个已安装TabNet的Jupyter笔记本。

新功能？

• 从版本**> 4.0**开始，注意力机制现在考虑了嵌入信息。这旨在即使在大量嵌入的情况下也能保持良好的注意力机制。此外，现在还可以指定注意力组（使用grouped_features）。注意力现在是在组级别而不是特征级别上进行的。如果数据集有许多来自单一数据源的列（例如：使用TF-IDF转换的文本列），这一点尤其有用。

贡献

在为TabNet仓库做出贡献时，请务必先通过新的或现有的问题讨论您希望进行的更改。

我们的提交遵循这里介绍的规则。

PyTorch TabNet能解决哪些问题？

• TabNetClassifier：二分类和多分类问题
• TabNetRegressor：简单回归和多任务回归问题
• TabNetMultiTaskClassifier：多任务多分类问题

如何使用？

TabNet现在与scikit-learn兼容，训练TabNetClassifier或TabNetRegressor非常容易。

from pytorch_tabnet.tab_model import TabNetClassifier, TabNetRegressor

clf = TabNetClassifier()  #TabNetRegressor()
clf.fit(
  X_train, Y_train,
  eval_set=[(X_valid, y_valid)]
)
preds = clf.predict(X_test)

或者对于TabNetMultiTaskClassifier：

from pytorch_tabnet.multitask import TabNetMultiTaskClassifier
clf = TabNetMultiTaskClassifier()
clf.fit(
  X_train, Y_train,
  eval_set=[(X_valid, y_valid)]
)
preds = clf.predict(X_test)

y_train/y_valid中的目标应包含单一类型（例如，它们必须全部是字符串或整数）。

默认评估指标

实现了一些经典的评估指标（自定义指标见下文）：

• 二分类指标：'auc'、'accuracy'、'balanced_accuracy'、'logloss'
• 多分类：'accuracy'、'balanced_accuracy'、'logloss'
• 回归：'mse'、'mae'、'rmse'、'rmsle'

重要提示：由于模型可能会预测负值，'rmsle'会自动将负预测截断为0。为了匹配给定的分数，在进行预测时需要使用np.clip(clf.predict(X_predict), a_min=0, a_max=None)。

自定义评估指标

您可以根据自己的需求创建指标。以下是一个基尼系数的示例（请注意，您需要指定该指标是应该最大化还是最小化）：

from pytorch_tabnet.metrics import Metric
from sklearn.metrics import roc_auc_score

class Gini(Metric):
    def __init__(self):
        self._name = "gini"
        self._maximize = True

    def __call__(self, y_true, y_score):
        auc = roc_auc_score(y_true, y_score[:, 1])
        return max(2*auc - 1, 0.)

clf = TabNetClassifier()
clf.fit(
  X_train, Y_train,
  eval_set=[(X_valid, y_valid)],
  eval_metric=[Gini]
)

一个专门的自定义示例笔记本在这里：https://github.com/dreamquark-ai/tabnet/blob/develop/customizing_example.ipynb

半监督预训练

半监督预训练后来被添加到TabNet的原始论文中，现在可以通过TabNetPretrainer类实现：

# TabNet预训练器
无监督模型 = TabNetPretrainer(
    优化器函数=torch.optim.Adam,
    优化器参数=dict(lr=2e-2),
    掩码类型='entmax' # "sparsemax"
)

无监督模型.fit(
    X_train=X_train,
    eval_set=[X_valid],
    预训练比例=0.8,
)

分类器 = TabNetClassifier(
    优化器函数=torch.optim.Adam,
    优化器参数=dict(lr=2e-2),
    调度器参数={"step_size":10, # 如何使用学习率调度器
                  "gamma":0.9},
    调度器函数=torch.optim.lr_scheduler.StepLR,
    掩码类型='sparsemax' # 如果使用预训练模型，此设置将被覆盖
)

分类器.fit(
    X_train=X_train, y_train=y_train,
    eval_set=[(X_train, y_train), (X_valid, y_valid)],
    eval_name=['train', 'valid'],
    eval_metric=['auc'],
    from_unsupervised=无监督模型
)

损失函数已被归一化，使其与pretraining_ratio、batch_size以及问题中特征的数量无关。 自监督损失大于1意味着模型的重建效果比单纯预测每个特征的均值还要差；而损失小于1则表示模型的表现优于仅预测均值。

完整的示例可以在笔记本文件pretraining_example.ipynb中找到。

/!\：当前实现试图重建原始输入，但批量归一化会应用一种无法通过单一线性变换推导出的随机变换，这使得重建变得更加困难。降低batch_size可能会使预训练更容易进行。

实时数据增强

现在可以在训练过程中应用自定义的数据增强流程。分类SMOTE和回归SMOTE的模板已添加到pytorch-tabnet/augmentations.py中，可以直接使用。

简便的保存与加载

保存和重新加载训练好的模型非常容易，这使得TabNet可以用于生产环境。

# 保存TabNet模型
saving_path_name = "./tabnet_model_test_1"
saved_filepath = clf.save_model(saving_path_name)

# 定义一个具有基本参数的新模型，并加载状态字典权重
loaded_clf = TabNetClassifier()
loaded_clf.load_model(saved_filepath)

有用链接

• 解释视频
• 二分类示例
• 多分类示例
• 回归示例
• 多任务回归示例
• 多任务多分类示例
• 使用TabNet的Kaggle MOA竞赛第一名解决方案

模型参数

• n_d : 整数（默认=8）

  决策预测层的宽度。较大的值会增加模型的容量，但也存在过拟合的风险。通常取值范围为8到64。

• n_a: 整数（默认=8）

  每个掩码的注意力嵌入宽度。 根据论文，n_d=n_a通常是不错的选择。（默认=8）

• n_steps : 整数（默认=3）

  架构中的步骤数（通常在3到10之间）

• gamma : 浮点数 （默认=1.3）

  这是掩码中特征重用的系数。 值接近1时，各层之间的掩码选择相关性最低。 取值范围为1.0到2.0。

• cat_idxs : 整数列表（默认=[] - 嵌入必需）

  分类特征索引列表。

• cat_dims : 整数列表（默认=[] - 嵌入必需）

  分类特征的类别数（即分类特征的唯一值数量） /!\ 不能预测新的类别

• cat_emb_dim : 整数列表（可选）

  每个分类特征的嵌入维度。（默认=1）

• n_independent : 整数 （默认=2）

  每个步骤中独立门控线性单元层的数量。 通常取值范围为1到5。

• n_shared : 整数（默认=2）

  每个步骤中共享的门控线性单元层数 通常取值范围为1到5

• epsilon : 浮点数 （默认 1e-15）

  应保持不变。

• seed : 整数（默认=0）

  用于复现性的随机种子

• momentum : 浮点数

  批量归一化的动量，通常在0.01到0.4之间（默认=0.02）

• clip_value : 浮点数（默认 None）

  如果给出浮点数值，则会将梯度裁剪至该值。

• lambda_sparse : 浮点数（默认 = 1e-3）

  这是原始论文中提出的额外稀疏性损失系数。该系数越大，模型在特征选择方面的稀疏性就越强。根据问题的难度，适当降低此值可能会有所帮助。

• optimizer_fn : torch.optim（默认=torch.optim.Adam）

  PyTorch优化器函数

• optimizer_params: 字典（默认=dict(lr=2e-2)）

  与optimizer_fn兼容的参数，用于初始化优化器。由于我们默认使用Adam优化器，因此这里用来定义训练时的初始学习率。正如原始论文所提到的，采用较大初始学习率0.02并配合衰减是一个不错的选择。

• scheduler_fn : torch.optim.lr_scheduler（默认=None）

  PyTorch学习率调度器，用于在训练过程中调整学习率。

• scheduler_params : 字典

  应用于调度器函数的参数字典。例如：{"gamma": 0.95, "step_size": 10}

• model_name : 字符串（默认 = 'DreamQuarkTabNet'）

  用于在磁盘上保存模型时的名称，您可以自定义此名称以便于检索和重用训练好的模型。

• verbose : 整数（默认=1）

  用于笔记本绘图的详细程度，设置为1可查看每一轮迭代，设置为0则不显示任何内容。

• device_name : 字符串（默认='auto'） 'cpu'表示在CPU上训练，'gpu'表示在GPU上训练，'auto'则自动检测GPU。

• mask_type: str（默认='sparsemax'） 可以是“sparsemax”或“entmax”：这是用于特征选择的掩码函数。

• grouped_features: 列表的列表的整数（默认=None） 这允许模型在同一组内的特征之间共享注意力。 当您的预处理生成了相关或依赖的特征时，这一点尤其有用：例如，您对文本列使用TF-IDF或PCA时。 请注意，在同一组内的特征之间，其重要性将是完全相同的。 此外，请注意，为分类变量生成的嵌入始终位于同一组内。

• n_shared_decoder : 整数（默认=1）

  解码器中共享GLU块的数量，这仅对TabNetPretrainer有用。

• n_indep_decoder : 整数（默认=1）

  解码器中独立GLU块的数量，这仅对TabNetPretrainer有用。

拟合参数

• X_train : np.array 或 scipy.sparse.csr_matrix

  训练特征

• y_train : np.array

  训练目标

• eval_set: 元组列表

  评估元组集合列表 (X, y)。
  最后一个用于早停

• eval_name: 字符串列表
  评估集名称列表。

• eval_metric : 字符串列表
  评估指标列表。
  最后一个指标用于早停。

• max_epochs : int (默认 = 200)

  训练的最大轮数。

• patience : int (默认 = 10)

  在执行早停之前，连续无提升的轮数。

  如果将耐心值设置为 0，则不会进行早停。

  注意，如果启用了耐心机制，那么在 fit 结束时会自动加载最佳轮次的权重。

• weights : int 或字典 (默认=0)

  /!\ 仅适用于 TabNetClassifier 抽样参数 0 : 不抽样 1 : 根据类别出现频率自动抽样 字典 : 键为类别，值为每个类别的权重

• loss_fn : torch.loss 或 torch.loss 列表

  训练用损失函数（回归默认为均方误差，分类默认为交叉熵）。 使用 TabNetMultiTaskClassifier 时，可以设置一个与任务数量相同长度的列表， 每个任务将分配到其专属的损失函数。

• batch_size : int (默认=1024)

  每个批次的样本数。建议使用较大的批次大小。

• virtual_batch_size : int (默认=128)

  用于“Ghost Batch Normalization”的小批次大小。 /!\ virtual_batch_size 应能整除 batch_size

• num_workers : int (默认=0)

  torch.utils.data.Dataloader 中使用的工作者数量。

• drop_last : bool (默认=False)

  训练过程中，如果最后一个批次不完整，是否将其丢弃。

• callbacks : 回调函数列表
  自定义回调函数列表。

• pretraining_ratio : 浮点数

  /!\ 仅限 TabNetPretrainer：预训练期间要掩码的输入特征百分比。 应介于 0 和 1 之间。数值越大，重建任务就越难。

• warm_start : bool (默认=False) 为了兼容 scikit-learn 的 API，该参数被设置为 False。 它允许对同一个模型进行两次拟合，并从热启动开始。

• compute_importance : bool (默认=True)

  是否计算特征重要性

TabNet 快速上手指南

TabNet 是一个基于 PyTorch 的可解释表格数据深度学习模型，具备类似 Scikit-learn 的简洁接口，支持分类、回归及多任务学习。

环境准备

• 操作系统：Linux, macOS, Windows
• Python 版本：3.7+
• 核心依赖：
  • PyTorch (需预先安装与硬件匹配的 CUDA 版本以启用 GPU 加速)
  • Scikit-learn
  • Pandas / Numpy
• 硬件建议：支持 CUDA 的 NVIDIA GPU 可显著提升训练速度（可选 CPU 模式）。

安装步骤

推荐使用 pip 进行安装。国内用户建议使用清华或阿里镜像源以加速下载。

方式一：使用 pip 安装（推荐）

# 使用默认源
pip install pytorch-tabnet

# 或使用国内镜像源（推荐中国开发者使用）
pip install pytorch-tabnet -i https://pypi.tuna.tsinghua.edu.cn/simple

方式二：使用 Conda 安装

conda install -c conda-forge pytorch-tabnet

源码安装（可选）

如需使用最新开发版或进行二次开发：

git clone git@github.com:dreamquark-ai/tabnet.git
cd tabnet
poetry install

基本使用

TabNet 的设计遵循 Scikit-learn 风格，只需几行代码即可完成模型定义、训练和预测。

1. 分类任务示例

适用于二分类和多分类问题。

from pytorch_tabnet.tab_model import TabNetClassifier

# 初始化模型
clf = TabNetClassifier(
    n_d=8, 
    n_a=8, 
    n_steps=3,
    gamma=1.3,
    cat_idxs=[],       # 填入类别特征索引，例如 [0, 2]
    cat_dims=[],       # 填入对应类别特征的维度（唯一值数量）
    cat_emb_dim=1,
    lambda_sparse=1e-3,
    optimizer_fn=torch.optim.Adam,
    optimizer_params=dict(lr=2e-2),
    scheduler_params={"step_size":10, "gamma":0.9},
    scheduler_fn=torch.optim.lr_scheduler.StepLR,
    mask_type='sparsemax',
    verbose=1,
    device_name='auto'
)

# 训练模型
# X_train, y_train: 训练集特征和标签 (Numpy array 或 Pandas DataFrame/Series)
# eval_set: 验证集列表，用于早停和监控
clf.fit(
    X_train=X_train, 
    y_train=y_train,
    eval_set=[(X_valid, y_valid)],
    eval_metric=['auc'] # 可选：'accuracy', 'logloss', 'balanced_accuracy' 等
)

# 预测
preds = clf.predict(X_test)
predict_proba = clf.predict_proba(X_test)

2. 回归任务示例

适用于单目标或多目标回归。

from pytorch_tabnet.tab_model import TabNetRegressor

reg = TabNetRegressor()

reg.fit(
    X_train=X_train, 
    y_train=y_train,
    eval_set=[(X_valid, y_valid)],
    eval_metric=['mse'] # 可选：'mae', 'rmse', 'rmsle'
)

preds = reg.predict(X_test)

3. 自定义评估指标

您可以轻松添加自定义指标（如 Gini 系数）：

from pytorch_tabnet.metrics import Metric
from sklearn.metrics import roc_auc_score

class Gini(Metric):
    def __init__(self):
        self._name = "gini"
        self._maximize = True

    def __call__(self, y_true, y_score):
        auc = roc_auc_score(y_true, y_score[:, 1])
        return max(2*auc - 1, 0.)

clf = TabNetClassifier()
clf.fit(
    X_train=X_train, 
    y_train=y_train,
    eval_set=[(X_valid, y_valid)],
    eval_metric=[Gini]
)

4. 模型保存与加载

TabNet 原生支持模型的持久化，便于生产部署。

# 保存模型
saving_path_name = "./tabnet_model_test_1"
saved_filepath = clf.save_model(saving_path_name)

# 加载模型
loaded_clf = TabNetClassifier()
loaded_clf.load_model(saved_filepath)

# 使用加载的模型预测
preds = loaded_clf.predict(X_test)

提示：对于包含大量类别的特征（如文本经 TF-IDF 处理后），建议在初始化时设置 grouped_features 参数以提升注意力机制的效果。