DOCSEARCHbyAlgolia
Submit Docs

tab-transformer-pytorch

repository·main·Indexed 22 days ago

https://github.com/lucidrains/tab-transformer-pytorch

A PyTorch implementation of the Tab Transformer and FT Transformer architectures for processing tabular data using attention mechanisms. Includes the TabTransformer class for categorical and continuous input tensors and the FTTransformer class for an alternative continuous value embedding scheme.

Tokens
1K
Snippets
3
Records
4
Agent score
28%

What's inside tab-transformer-pytorch

  1. Perform unsupervised training with Tab Transformer

    main
    To perform unsupervised training as described in the original paper, you can convert your category tokens to unique IDs and then apply the Electra algorithm (available in lucidrains/electra-pytorch) directly on the model.transformer component.

  3. Use the TabTransformer class

    main

    The TabTransformer class implements an attention network for tabular data. It requires both categorical and continuous input tensors.

    Constructor Arguments

    • categories: A tuple containing the number of unique values within each category.
    • num_continuous: The number of continuous (numerical) values.
    • dim: The embedding dimension (the paper uses 32).
    • dim_out: The output dimension (e.g., 1 for binary prediction).
    • depth: The number of transformer layers (the paper recommends 6).
    • heads: The number of attention heads (the paper recommends 8).
    • attn_dropout: Post-attention dropout rate.
    • ff_dropout: Feed forward dropout rate.
    • mlp_hidden_mults: A tuple of relative multiples of each hidden dimension of the last MLP leading to logits.
    • mlp_act: Activation function for the final MLP (defaults to nn.ReLU()).
    • continuous_mean_std: (Optional) A tensor used to normalize continuous values before layer norm.

    Forward Pass

    Pass x_categ (categorical values as integers) and x_cont (continuous values) to the model.

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

cont_mean_std = torch.randn(10, 2)

model = TabTransformer(
    categories = (10, 5, 6, 5, 8),
    num_continuous = 10,
    dim = 32,
    dim_out = 1,
    depth = 6,
    heads = 8,
    attn_dropout = 0.1,
    ff_dropout = 0.1,
    mlp_hidden_mults = (4, 2),
    mlp_act = nn.ReLU(),
    continuous_mean_std = cont_mean_std
)

x_categ = torch.randint(0, 5, (1, 5))
# assume continuous values are already normalized individually
x_cont = torch.randn(1, 10)

pred = model(x_categ, x_cont) # (1, 1)

  4. Use the FTTransformer class

    main

    The FTTransformer class is an improved version of the Tab Transformer that uses a simpler scheme for embedding continuous numerical values. It is included for comparison purposes.

    Constructor Arguments

    • categories: A tuple containing the number of unique values within each category.
    • num_continuous: The number of continuous (numerical) values.
    • dim: The embedding dimension.
    • dim_out: The output dimension.
    • depth: The number of transformer layers.
    • heads: The number of attention heads.
    • attn_dropout: Post-attention dropout rate.
    • ff_dropout: Feed forward dropout rate.

    Forward Pass

    Pass x_categ (categorical values as integers) and x_numer (numerical values) to the model.

    import torch
from tab_transformer_pytorch import FTTransformer

model = FTTransformer(
    categories = (10, 5, 6, 5, 8),
    num_continuous = 10,
    dim = 32,
    dim_out = 1,
    depth = 6,
    heads = 8,
    attn_dropout = 0.1,
    ff_dropout = 0.1
)

x_categ = torch.randint(0, 5, (1, 5))
x_numer = torch.randn(1, 10)

pred = model(x_categ, x_numer) # (1, 1)