pytorch-frame

一个用于在异构表格数据上构建神经网络模型的模块化深度学习框架。

---

文档 | 论文

PyTorch Frame 是 PyTorch 的深度学习扩展，专为包含数值、类别、时间、文本和图像等不同列类型的异构表格数据设计。它提供了一个模块化框架，用于实现现有和未来的方法。该库包含最先进模型的方法、用户友好的小批量加载器、基准数据集以及自定义数据集成接口。

PyTorch Frame 让表格数据的深度学习研究变得更加普及，既适合新手也适合专家。我们的目标是：

1. 促进表格数据的深度学习： 历史上，基于树的模型（如 GBDT）在表格学习方面表现出色，但存在一些明显的局限性，例如与下游模型的集成困难，以及处理复杂列类型（如文本、序列和嵌入）的能力。深度表格模型有望解决这些局限性。我们的目标是通过模块化实现并支持多样的列类型，来促进表格数据的深度学习研究。

2. 与大型语言模型等多样化模型架构集成： PyTorch Frame 支持与各种不同的架构集成，包括大型语言模型。使用任何下载的模型或嵌入 API 端点，您可以为文本数据生成嵌入，并与其他复杂语义类型一起使用深度学习模型进行训练。我们支持以下（但不限于）：

OpenAI 嵌入代码示例

Cohere Embed v3 代码示例

Hugging Face 代码示例

Voyage AI 代码示例

• 库亮点
• 架构概览
• 快速上手
• 已实现的深度表格模型
• 基准测试
• 安装

库亮点

PyTorch Frame 直接构建在 PyTorch 之上，确保现有 PyTorch 用户能够顺利过渡。主要特性包括：

• 多样化列类型： PyTorch Frame 支持跨各种列类型的学习：numerical、categorical、multicategorical、text_embedded、text_tokenized、timestamp、image_embedded 和 embedding。详细教程请参见此处。
• 模块化模型设计： 支持模块化深度学习模型实现，促进代码重用、清晰编码和实验灵活性。更多详情请参见架构概览。
• 模型 实现了许多最先进的深度表格模型以及强大的 GBDT（XGBoost、CatBoost 和 LightGBM），并支持超参数调优。
• 数据集： 提供了一系列可直接使用的表格数据集。同时支持自定义数据集以解决您自己的问题。 我们对深度表格模型与 GBDT 进行了基准测试。
• PyTorch 集成： 与其他 PyTorch 库无缝集成，便于 PyTorch Frame 与下游 PyTorch 模型的端到端训练。例如，通过与 PyTorch 图神经网络库 PyG 集成，我们可以对关系数据库进行深度学习。更多信息请参见 RelBench 和示例代码（进行中）。

架构概览

PyTorch Frame 中的模型遵循 FeatureEncoder、TableConv 和 Decoder 的模块化设计，如下图所示：

本质上，这种模块化设置使用户能够轻松尝试各种架构：

• Materialization 处理将原始 pandas DataFrame 转换为适合 PyTorch 训练和建模的 TensorFrame。
• FeatureEncoder 将 TensorFrame 编码为大小为 [batch_size, num_cols, channels] 的隐藏列嵌入。
• TableConv 对隐藏嵌入进行列间交互建模。
• Decoder 为每行生成嵌入/预测。

快速上手

在这个快速上手中，我们将展示如何仅用几行代码就能轻松创建和训练一个深度表格模型。

构建并训练您自己的深度表格模型

作为示例，我们将按照 PyTorch Frame 的模块化架构实现一个简单的 ExampleTransformer。 在下面的示例中：

• self.encoder 将输入的 TensorFrame 映射到大小为 [batch_size, num_cols, channels] 的嵌入。
• self.convs 迭代地将大小为 [batch_size, num_cols, channels] 的嵌入转换为相同大小的嵌入。
• self.decoder 将大小为 [batch_size, num_cols, channels] 的嵌入池化为 [batch_size, out_channels]。

from torch import Tensor
from torch.nn import Linear, Module, ModuleList

from torch_frame import TensorFrame, stype
from torch_frame.nn.conv import TabTransformerConv
from torch_frame.nn.encoder import (
    EmbeddingEncoder,
    LinearEncoder,
    StypeWiseFeatureEncoder,
)

class ExampleTransformer(Module):
    def __init__(
        self,
        channels, out_channels, num_layers, num_heads,
        col_stats, col_names_dict,
    ):
        super().__init__()
        self.encoder = StypeWiseFeatureEncoder(
            out_channels=channels,
            col_stats=col_stats,
            col_names_dict=col_names_dict,
            stype_encoder_dict={
                stype.categorical: EmbeddingEncoder(),
                stype.numerical: LinearEncoder()
            },
        )
        self.convs = ModuleList([
            TabTransformerConv(
                channels=channels,
                num_heads=num_heads,
            ) for _ in range(num_layers)
        ])
        self.decoder = Linear(channels, out_channels)

    def forward(self, tf: TensorFrame) -> Tensor:
        x, _ = self.encoder(tf)
        for conv in self.convs:
            x = conv(x)
        out = self.decoder(x.mean(dim=1))
        return out

要准备数据，我们可以快速实例化一个预定义的数据集并创建一个与 PyTorch 兼容的数据加载器，如下所示：

from torch_frame.datasets import Yandex
from torch_frame.data import DataLoader

dataset = Yandex(root='/tmp/adult', name='adult')
dataset.materialize()
train_dataset = dataset[:0.8]
train_loader = DataLoader(train_dataset.tensor_frame, batch_size=128,
                          shuffle=True)

然后，我们只需按照标准PyTorch训练流程来优化模型参数。就这么简单！

import torch
import torch.nn.functional as F

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = ExampleTransformer(
    channels=32,
    out_channels=dataset.num_classes,
    num_layers=2,
    num_heads=8,
    col_stats=train_dataset.col_stats,
    col_names_dict=train_dataset.tensor_frame.col_names_dict,
).to(device)

optimizer = torch.optim.Adam(model.parameters())

for epoch in range(50):
    for tf in train_loader:
        tf = tf.to(device)
        pred = model.forward(tf)
        loss = F.cross_entropy(pred, tf.y)
        optimizer.zero_grad()
        loss.backward()

已实现的深度表格模型

以下是目前支持的深度表格模型列表：

• Trompt，来自Chen等人：Trompt: 更优的表格数据深度神经网络（ICML 2023）[示例]
• FTTransformer，来自Gorishniy等人：重新审视表格数据的深度学习模型（NeurIPS 2021）[示例]
• ResNet，来自Gorishniy等人：重新审视表格数据的深度学习模型（NeurIPS 2021）[示例]
• TabNet，来自Arık等人：TabNet：注意力可解释的表格学习（AAAI 2021）[示例]
• ExcelFormer，来自Chen等人：ExcelFormer：一个在表格数据上胜过GBDT的神经网络 [示例]
• TabTransformer，来自Huang等人：TabTransformer：使用上下文嵌入的表格数据建模 [示例]

此外，我们还为想要将模型性能与GBDT进行比较的用户实现了XGBoost、CatBoost和LightGBM的示例，这些示例使用Optuna进行超参数调优。

基准测试

我们在多种规模和任务类型的公共数据集上对最近的表格深度学习模型与GBDT进行了基准测试。

下图展示了各种模型在小型回归数据集上的性能，其中行代表模型名称，列代表数据集索引（这里我们有13个数据集）。有关分类和更大数据集的更多结果，请查看基准测试文档。

我们可以看到，一些最新的深度表格模型能够达到与强大的GBDT相当的模型性能（尽管训练速度慢5-100倍）。使深度表格模型在更少计算资源下表现更好是未来研究的一个富有成果的方向。

我们还在一个带有一列文本的真实世界表格数据集（葡萄酒评论）上对不同的文本编码器进行了基准测试。下表显示了性能结果：

Hugging Face文本编码器的基准测试脚本在这个文件中，其他文本编码器的基准测试脚本在这个文件中。

安装

PyTorch Frame支持Python 3.8到Python 3.11版本。

pip install pytorch_frame

查看安装指南了解其他安装选项。

引用

如果您在工作中使用了PyTorch Frame，请引用我们的论文（Bibtex如下）。

@article{hu2024pytorch,
  title={PyTorch Frame: A Modular Framework for Multi-Modal Tabular Learning},
  author={Hu, Weihua and Yuan, Yiwen and Zhang, Zecheng and Nitta, Akihiro and Cao, Kaidi and Kocijan, Vid and Leskovec, Jure and Fey, Matthias},
  journal={arXiv preprint arXiv:2404.00776},
  year={2024}
}