【免费下载】 TabTransformer-PyTorch 使用教程

2026-01-16 10:32:42作者：伍希望

1. 项目介绍

TabTransformer-PyTorch 是一个基于PyTorch实现的用于处理表格数据的注意力网络模型。这个开源项目由 Lucidrains 创建，它提供了一个简单架构来利用Transformer的注意力机制进行特征学习，几乎可以达到梯度提升决策树（GBDT）在表结构数据预测任务上的性能。根据作者更新的信息，亚马逊AI已经在实际场景中用注意力机制在表格数据上超越了GBDT。

2. 项目快速启动

首先确保已安装PyTorch，然后通过以下命令安装tab-transformer-pytorch库：

pip install tab-transformer-pytorch

接下来是一个简单的示例，展示如何创建和初始化TabTransformer模型：

import torch
import torch.nn as nn
from tab_transformer_pytorch import TabTransformer

# 假设我们有10个分类特征，每个有5个唯一值；接着是5个分类特征，每个有6个唯一值；
# 然后是5个分类特征，每个有5个唯一值；另外还有8个连续特征。
categories = (10, 5, 6, 5, 8)
num_continuous = 10

# 配置模型参数
model = TabTransformer(
    categories=categories,
    num_continuous=num_continuous,
    dim=32,  # 默认为32维度
    dim_out=1,  # 输出层维度，例如二分类问题设为1
    depth=6,  # 模型深度，推荐6层
    heads=8,  # 多头注意力的头数，推荐8
    attn_dropout=0.1,  # 注意力层的Dropout比例
    ff_dropout=0.1,  # FFN层的Dropout比例
    mlp_hidden_mults=(4, 2),  # 最后一层到logits层MLP的隐藏维度相对倍数
    mlp_act=nn.ReLU(),  # 最后一层MLP的激活函数，默认ReLU
)

现在你可以将数据输入模型进行前向传播了。

3. 应用案例和最佳实践

3.1 数据预处理

在使用TabTransformer时，务必对你的表格数据进行适当的预处理，包括标准化连续数值特征以及编码分类特征。

# 示例数据预处理
cont_data = torch.randn(numContinuous, batch_size)
cat_data = ...  # 对应于每个类别特征的编码数据

# 将连续数据归一化
cont_mean_std = torch.randn(num_continuous, 2)
cont_normalized = (cont_data - cont_mean_std[0]) / cont_mean_std[1]

# 拼接分类和连续数据
input_data = torch.cat([cat_data, cont_normalized], dim=-1)

3.2 训练和评估

训练时，记得设置优化器并定义损失函数，以适应你的具体任务（如二元分类或回归）。在每个epoch之后，使用验证集计算性能指标，例如准确率、AUC等。

optimizer = torch.optim.Adam(model.parameters())
criterion = nn.BCEWithLogitsLoss()

for epoch in range(num_epochs):
    model.train()
    for inputs, labels in train_dataloader:
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
    
    # 在验证集上评估
    model.eval()
    with torch.no_grad():
        val_loss, val_metric = evaluate_model(model, val_dataloader, criterion, metric_function)
    print(f'Epoch {epoch+1}/{num_epochs}: Loss={val_loss}, Metric={val_metric}')