title structured data learning with tabtransformer khalid salamahttps www linkedin cominkhalidsalama24403144 date created 20220118 last modified 20220118 description using contextual embeddings for structured data classification accelerator gpu introduction this example demonstrates how to do structured data classification using tabtransformerhttps arxiv orgabs2012 06678 a deep tabular data modeling architecture for supervised and semisupervised learning the tabtransformer is built upon selfattention based transformers the transformer layers transform the embeddings of categorical features into robust contextual embeddings to achieve higher predictive accuracy setup prepare the data this example uses the united states census income datasethttps archive ics uci edumldatasetscensusincome provided by the uc irvine machine learning repositoryhttps archive ics uci edumlindex php the task is binary classification to predict whether a person is likely to be making over usd 50 000 a year the dataset includes 48 842 instances with 14 input features 5 numerical features and 9 categorical features first let s load the dataset from the uci machine learning repository into a pandas dataframe remove the first record because it is not a valid data example and a trailing dot in the class labels now we store the training and test data in separate csv files define dataset metadata here we define the metadata of the dataset that will be useful for reading and parsing the data into input features and encoding the input features with respect to their types a list of the numerical feature names a dictionary of the categorical features and their vocabulary name of the column to be used as instances weight a list of the categorical feature names a list of all the input features a list of column default values for each feature the name of the target feature a list of the labels of the target features configure the hyperparameters the hyperparameters includes model architecture and training configurations implement data reading pipeline we define an input function that reads and parses the file then converts features and labels into atf data datasethttps www tensorflow orgguidedatasets for training or evaluation create a lookup to convert a string values to an integer indices since we are not using a mask token nor expecting any out of vocabulary oov token we set masktoken to none and numoovindices to 0 implement a training and evaluation procedure create model inputs now define the inputs for the models as a dictionary where the key is the feature name and the value is a keras layers input tensor with the corresponding feature shape and data type encode features the encodeinputs method returns encodedcategoricalfeaturelist and numericalfeaturelist we encode the categorical features as embeddings using a fixed embeddingdims for all the features regardless their vocabulary sizes this is required for the transformer model create a lookup to convert a string values to an integer indices since we are not using a mask token nor expecting any out of vocabulary oov token we set masktoken to none and numoovindices to 0 convert the string input values into integer indices create an embedding layer with the specified dimensions convert the index values to embedding representations use the numerical features asis implement an mlp block experiment 1 a baseline model in the first experiment we create a simple multilayer feedforward network create model inputs encode features concatenate all features compute feedforward layer units create several feedforwad layers with skip connections compute mlp hiddenunits create final mlp add a sigmoid as a binary classifer let s train and evaluate the baseline model the baseline linear model achieves 81 validation accuracy experiment 2 tabtransformer the tabtransformer architecture works as follows 1 all the categorical features are encoded as embeddings using the same embeddingdims this means that each value in each categorical feature will have its own embedding vector 2 a column embedding one embedding vector for each categorical feature is added pointwise to the categorical feature embedding 3 the embedded categorical features are fed into a stack of transformer blocks each transformer block consists of a multihead selfattention layer followed by a feedforward layer 3 the outputs of the final transformer layer which are the contextual embeddings of the categorical features are concatenated with the input numerical features and fed into a final mlp block 4 a softmax classifer is applied at the end of the model the paperhttps arxiv orgabs2012 06678 discusses both addition and concatenation of the column embedding in the appendix experiment and model details section the architecture of tabtransformer is shown below as presented in the paper img srchttps raw githubusercontent comkerasteamkerasiomasterexamplesstructureddataimgtabtransformertabtransformer png width500 create model inputs encode features stack categorical feature embeddings for the tansformer concatenate numerical features add column embedding to categorical feature embeddings create multiple layers of the transformer block create a multihead attention layer skip connection 1 layer normalization 1 feedforward skip connection 2 layer normalization 2 flatten the contextualized embeddings of the categorical features apply layer normalization to the numerical features prepare the input for the final mlp block compute mlp hiddenunits create final mlp add a sigmoid as a binary classifer let s train and evaluate the tabtransformer model the tabtransformer model achieves 85 validation accuracy note that with the default parameter configurations both the baseline and the tabtransformer have similar number of trainable weights 109 629 and 92 151 respectively and both use the same training hyperparameters conclusion tabtransformer significantly outperforms mlp and recent deep networks for tabular data while matching the performance of treebased ensemble models tabtransformer can be learned in endtoend supervised training using labeled examples for a scenario where there are a few labeled examples and a large number of unlabeled examples a pretraining procedure can be employed to train the transformer layers using unlabeled data this is followed by finetuning of the pretrained transformer layers along with the top mlp layer using the labeled data example available on huggingface trained model demo generic badgehttps img shields iobadge20modeltabtransformerblack svghttps huggingface cokerasiotabtransformer generic badgehttps img shields iobadge20spacestabtransformerblack svghttps huggingface cospaceskerasiotabtransformerclassification introduction this example demonstrates how to do structured data classification using tabtransformer https arxiv org abs 2012 06678 a deep tabular data modeling architecture for supervised and semi supervised learning the tabtransformer is built upon self attention based transformers the transformer layers transform the embeddings of categorical features into robust contextual embeddings to achieve higher predictive accuracy setup prepare the data this example uses the united states census income dataset https archive ics uci edu ml datasets census income provided by the uc irvine machine learning repository https archive ics uci edu ml index php the task is binary classification to predict whether a person is likely to be making over usd 50 000 a year the dataset includes 48 842 instances with 14 input features 5 numerical features and 9 categorical features first let s load the dataset from the uci machine learning repository into a pandas dataframe remove the first record because it is not a valid data example and a trailing dot in the class labels now we store the training and test data in separate csv files define dataset metadata here we define the metadata of the dataset that will be useful for reading and parsing the data into input features and encoding the input features with respect to their types a list of the numerical feature names a dictionary of the categorical features and their vocabulary name of the column to be used as instances weight a list of the categorical feature names a list of all the input features a list of column default values for each feature the name of the target feature a list of the labels of the target features configure the hyperparameters the hyperparameters includes model architecture and training configurations number of transformer blocks number of attention heads embedding dimensions of the categorical features mlp hidden layer units as factors of the number of inputs number of mlp blocks in the baseline model implement data reading pipeline we define an input function that reads and parses the file then converts features and labels into a tf data dataset https www tensorflow org guide datasets for training or evaluation create a lookup to convert a string values to an integer indices since we are not using a mask token nor expecting any out of vocabulary oov token we set mask_token to none and num_oov_indices to 0 implement a training and evaluation procedure create model inputs now define the inputs for the models as a dictionary where the key is the feature name and the value is a keras layers input tensor with the corresponding feature shape and data type encode features the encode_inputs method returns encoded_categorical_feature_list and numerical_feature_list we encode the categorical features as embeddings using a fixed embedding_dims for all the features regardless their vocabulary sizes this is required for the transformer model create a lookup to convert a string values to an integer indices since we are not using a mask token nor expecting any out of vocabulary oov token we set mask_token to none and num_oov_indices to 0 convert the string input values into integer indices create an embedding layer with the specified dimensions convert the index values to embedding representations use the numerical features as is implement an mlp block experiment 1 a baseline model in the first experiment we create a simple multi layer feed forward network create model inputs encode features concatenate all features compute feedforward layer units create several feedforwad layers with skip connections compute mlp hidden_units create final mlp add a sigmoid as a binary classifer let s train and evaluate the baseline model the baseline linear model achieves 81 validation accuracy experiment 2 tabtransformer the tabtransformer architecture works as follows 1 all the categorical features are encoded as embeddings using the same embedding_dims this means that each value in each categorical feature will have its own embedding vector 2 a column embedding one embedding vector for each categorical feature is added point wise to the categorical feature embedding 3 the embedded categorical features are fed into a stack of transformer blocks each transformer block consists of a multi head self attention layer followed by a feed forward layer 3 the outputs of the final transformer layer which are the contextual embeddings of the categorical features are concatenated with the input numerical features and fed into a final mlp block 4 a softmax classifer is applied at the end of the model the paper https arxiv org abs 2012 06678 discusses both addition and concatenation of the column embedding in the appendix experiment and model details section the architecture of tabtransformer is shown below as presented in the paper img src https raw githubusercontent com keras team keras io master examples structured_data img tabtransformer tabtransformer png width 500 create model inputs encode features stack categorical feature embeddings for the tansformer concatenate numerical features add column embedding to categorical feature embeddings create multiple layers of the transformer block create a multi head attention layer skip connection 1 layer normalization 1 feedforward using partial to provide keyword arguments before initialization skip connection 2 layer normalization 2 flatten the contextualized embeddings of the categorical features apply layer normalization to the numerical features prepare the input for the final mlp block compute mlp hidden_units create final mlp add a sigmoid as a binary classifer let s train and evaluate the tabtransformer model the tabtransformer model achieves 85 validation accuracy note that with the default parameter configurations both the baseline and the tabtransformer have similar number of trainable weights 109 629 and 92 151 respectively and both use the same training hyperparameters conclusion tabtransformer significantly outperforms mlp and recent deep networks for tabular data while matching the performance of tree based ensemble models tabtransformer can be learned in end to end supervised training using labeled examples for a scenario where there are a few labeled examples and a large number of unlabeled examples a pre training procedure can be employed to train the transformer layers using unlabeled data this is followed by fine tuning of the pre trained transformer layers along with the top mlp layer using the labeled data example available on huggingface trained model demo generic badge https img shields io badge 20model tabtransformer black svg https huggingface co keras io tab_transformer generic badge https img shields io badge 20spaces tabtransformer black svg https huggingface co spaces keras io tabtransformer_classification
import keras
from keras import layers
from keras import ops

import math
import numpy as np
import pandas as pd
from tensorflow import data as tf_data
import matplotlib.pyplot as plt
from functools import partial



CSV_HEADER = [
    "age",
    "workclass",
    "fnlwgt",
    "education",
    "education_num",
    "marital_status",
    "occupation",
    "relationship",
    "race",
    "gender",
    "capital_gain",
    "capital_loss",
    "hours_per_week",
    "native_country",
    "income_bracket",
]

train_data_url = (
    "https://archive.ics.uci.edu/ml/machine-learning-databases/adult/adult.data"
)
train_data = pd.read_csv(train_data_url, header=None, names=CSV_HEADER)

test_data_url = (
    "https://archive.ics.uci.edu/ml/machine-learning-databases/adult/adult.test"
)
test_data = pd.read_csv(test_data_url, header=None, names=CSV_HEADER)

print(f"Train dataset shape: {train_data.shape}")
print(f"Test dataset shape: {test_data.shape}")



test_data = test_data[1:]
test_data.income_bracket = test_data.income_bracket.apply(
    lambda value: value.replace(".", "")
)



train_data_file = "train_data.csv"
test_data_file = "test_data.csv"

train_data.to_csv(train_data_file, index=False, header=False)
test_data.to_csv(test_data_file, index=False, header=False)




NUMERIC_FEATURE_NAMES = [
    "age",
    "education_num",
    "capital_gain",
    "capital_loss",
    "hours_per_week",
]

CATEGORICAL_FEATURES_WITH_VOCABULARY = {
    "workclass": sorted(list(train_data["workclass"].unique())),
    "education": sorted(list(train_data["education"].unique())),
    "marital_status": sorted(list(train_data["marital_status"].unique())),
    "occupation": sorted(list(train_data["occupation"].unique())),
    "relationship": sorted(list(train_data["relationship"].unique())),
    "race": sorted(list(train_data["race"].unique())),
    "gender": sorted(list(train_data["gender"].unique())),
    "native_country": sorted(list(train_data["native_country"].unique())),
}

WEIGHT_COLUMN_NAME = "fnlwgt"

CATEGORICAL_FEATURE_NAMES = list(CATEGORICAL_FEATURES_WITH_VOCABULARY.keys())

FEATURE_NAMES = NUMERIC_FEATURE_NAMES + CATEGORICAL_FEATURE_NAMES

COLUMN_DEFAULTS = [
    [0.0] if feature_name in NUMERIC_FEATURE_NAMES + [WEIGHT_COLUMN_NAME] else ["NA"]
    for feature_name in CSV_HEADER
]

TARGET_FEATURE_NAME = "income_bracket"

TARGET_LABELS = [" <=50K", " >50K"]



LEARNING_RATE = 0.001
WEIGHT_DECAY = 0.0001
DROPOUT_RATE = 0.2
BATCH_SIZE = 265
NUM_EPOCHS = 15

NUM_TRANSFORMER_BLOCKS = 3  
NUM_HEADS = 4  
EMBEDDING_DIMS = 16  
MLP_HIDDEN_UNITS_FACTORS = [
    2,
    1,
]  
NUM_MLP_BLOCKS = 2  



target_label_lookup = layers.StringLookup(
    vocabulary=TARGET_LABELS, mask_token=None, num_oov_indices=0
)


def prepare_example(features, target):
    target_index = target_label_lookup(target)
    weights = features.pop(WEIGHT_COLUMN_NAME)
    return features, target_index, weights


lookup_dict = {}
for feature_name in CATEGORICAL_FEATURE_NAMES:
    vocabulary = CATEGORICAL_FEATURES_WITH_VOCABULARY[feature_name]
    
    
    
    lookup = layers.StringLookup(
        vocabulary=vocabulary, mask_token=None, num_oov_indices=0
    )
    lookup_dict[feature_name] = lookup


def encode_categorical(batch_x, batch_y, weights):
    for feature_name in CATEGORICAL_FEATURE_NAMES:
        batch_x[feature_name] = lookup_dict[feature_name](batch_x[feature_name])

    return batch_x, batch_y, weights


def get_dataset_from_csv(csv_file_path, batch_size=128, shuffle=False):
    dataset = (
        tf.data.experimental.make_csv_dataset(
            csv_file_path,
            batch_size=batch_size,
            column_names=CSV_HEADER,
            column_defaults=COLUMN_DEFAULTS,
            label_name=TARGET_FEATURE_NAME,
            num_epochs=1,
            header=False,
            na_value="?",
            shuffle=shuffle,
        )
        .map(prepare_example, num_parallel_calls=tf.data.AUTOTUNE, deterministic=False)
        .map(encode_categorical)
    )
    return dataset.cache()





def run_experiment(
    model,
    train_data_file,
    test_data_file,
    num_epochs,
    learning_rate,
    weight_decay,
    batch_size,
):
    optimizer = keras.optimizers.AdamW(
        learning_rate=learning_rate, weight_decay=weight_decay
    )

    model.compile(
        optimizer=optimizer,
        loss=keras.losses.BinaryCrossentropy(),
        metrics=[keras.metrics.BinaryAccuracy(name="accuracy")],
    )

    train_dataset = get_dataset_from_csv(train_data_file, batch_size, shuffle=True)
    validation_dataset = get_dataset_from_csv(test_data_file, batch_size)

    print("Start training the model...")
    history = model.fit(
        train_dataset, epochs=num_epochs, validation_data=validation_dataset
    )
    print("Model training finished")

    _, accuracy = model.evaluate(validation_dataset, verbose=0)

    print(f"Validation accuracy: {round(accuracy * 100, 2)}%")

    return history





def create_model_inputs():
    inputs = {}
    for feature_name in FEATURE_NAMES:
        if feature_name in NUMERIC_FEATURE_NAMES:
            inputs[feature_name] = layers.Input(
                name=feature_name, shape=(), dtype="float32"
            )
        else:
            inputs[feature_name] = layers.Input(
                name=feature_name, shape=(), dtype="float32"
            )
    return inputs





def encode_inputs(inputs, embedding_dims):
    encoded_categorical_feature_list = []
    numerical_feature_list = []

    for feature_name in inputs:
        if feature_name in CATEGORICAL_FEATURE_NAMES:
            vocabulary = CATEGORICAL_FEATURES_WITH_VOCABULARY[feature_name]
            
            
            

            

            
            embedding = layers.Embedding(
                input_dim=len(vocabulary), output_dim=embedding_dims
            )

            
            encoded_categorical_feature = embedding(inputs[feature_name])
            encoded_categorical_feature_list.append(encoded_categorical_feature)

        else:
            
            numerical_feature = ops.expand_dims(inputs[feature_name], -1)
            numerical_feature_list.append(numerical_feature)

    return encoded_categorical_feature_list, numerical_feature_list





def create_mlp(hidden_units, dropout_rate, activation, normalization_layer, name=None):
    mlp_layers = []
    for units in hidden_units:
        mlp_layers.append(normalization_layer()),
        mlp_layers.append(layers.Dense(units, activation=activation))
        mlp_layers.append(layers.Dropout(dropout_rate))

    return keras.Sequential(mlp_layers, name=name)





def create_baseline_model(
    embedding_dims, num_mlp_blocks, mlp_hidden_units_factors, dropout_rate
):
    
    inputs = create_model_inputs()
    
    encoded_categorical_feature_list, numerical_feature_list = encode_inputs(
        inputs, embedding_dims
    )
    
    features = layers.concatenate(
        encoded_categorical_feature_list + numerical_feature_list
    )
    
    feedforward_units = [features.shape[-1]]

    
    for layer_idx in range(num_mlp_blocks):
        features = create_mlp(
            hidden_units=feedforward_units,
            dropout_rate=dropout_rate,
            activation=keras.activations.gelu,
            normalization_layer=layers.LayerNormalization,
            name=f"feedforward_{layer_idx}",
        )(features)

    
    mlp_hidden_units = [
        factor * features.shape[-1] for factor in mlp_hidden_units_factors
    ]
    
    features = create_mlp(
        hidden_units=mlp_hidden_units,
        dropout_rate=dropout_rate,
        activation=keras.activations.selu,
        normalization_layer=layers.BatchNormalization,
        name="MLP",
    )(features)

    
    outputs = layers.Dense(units=1, activation="sigmoid", name="sigmoid")(features)
    model = keras.Model(inputs=inputs, outputs=outputs)
    return model


baseline_model = create_baseline_model(
    embedding_dims=EMBEDDING_DIMS,
    num_mlp_blocks=NUM_MLP_BLOCKS,
    mlp_hidden_units_factors=MLP_HIDDEN_UNITS_FACTORS,
    dropout_rate=DROPOUT_RATE,
)

print("Total model weights:", baseline_model.count_params())
keras.utils.plot_model(baseline_model, show_shapes=True, rankdir="LR")



history = run_experiment(
    model=baseline_model,
    train_data_file=train_data_file,
    test_data_file=test_data_file,
    num_epochs=NUM_EPOCHS,
    learning_rate=LEARNING_RATE,
    weight_decay=WEIGHT_DECAY,
    batch_size=BATCH_SIZE,
)






def create_tabtransformer_classifier(
    num_transformer_blocks,
    num_heads,
    embedding_dims,
    mlp_hidden_units_factors,
    dropout_rate,
    use_column_embedding=False,
):
    
    inputs = create_model_inputs()
    
    encoded_categorical_feature_list, numerical_feature_list = encode_inputs(
        inputs, embedding_dims
    )
    
    encoded_categorical_features = ops.stack(encoded_categorical_feature_list, axis=1)
    
    numerical_features = layers.concatenate(numerical_feature_list)

    
    if use_column_embedding:
        num_columns = encoded_categorical_features.shape[1]
        column_embedding = layers.Embedding(
            input_dim=num_columns, output_dim=embedding_dims
        )
        column_indices = ops.arange(start=0, stop=num_columns, step=1)
        encoded_categorical_features = encoded_categorical_features + column_embedding(
            column_indices
        )

    
    for block_idx in range(num_transformer_blocks):
        
        attention_output = layers.MultiHeadAttention(
            num_heads=num_heads,
            key_dim=embedding_dims,
            dropout=dropout_rate,
            name=f"multihead_attention_{block_idx}",
        )(encoded_categorical_features, encoded_categorical_features)
        
        x = layers.Add(name=f"skip_connection1_{block_idx}")(
            [attention_output, encoded_categorical_features]
        )
        
        x = layers.LayerNormalization(name=f"layer_norm1_{block_idx}", epsilon=1e-6)(x)
        
        feedforward_output = create_mlp(
            hidden_units=[embedding_dims],
            dropout_rate=dropout_rate,
            activation=keras.activations.gelu,
            normalization_layer=partial(
                layers.LayerNormalization, epsilon=1e-6
            ),  
            name=f"feedforward_{block_idx}",
        )(x)
        
        x = layers.Add(name=f"skip_connection2_{block_idx}")([feedforward_output, x])
        
        encoded_categorical_features = layers.LayerNormalization(
            name=f"layer_norm2_{block_idx}", epsilon=1e-6
        )(x)

    
    categorical_features = layers.Flatten()(encoded_categorical_features)
    
    numerical_features = layers.LayerNormalization(epsilon=1e-6)(numerical_features)
    
    features = layers.concatenate([categorical_features, numerical_features])

    
    mlp_hidden_units = [
        factor * features.shape[-1] for factor in mlp_hidden_units_factors
    ]
    
    features = create_mlp(
        hidden_units=mlp_hidden_units,
        dropout_rate=dropout_rate,
        activation=keras.activations.selu,
        normalization_layer=layers.BatchNormalization,
        name="MLP",
    )(features)

    
    outputs = layers.Dense(units=1, activation="sigmoid", name="sigmoid")(features)
    model = keras.Model(inputs=inputs, outputs=outputs)
    return model


tabtransformer_model = create_tabtransformer_classifier(
    num_transformer_blocks=NUM_TRANSFORMER_BLOCKS,
    num_heads=NUM_HEADS,
    embedding_dims=EMBEDDING_DIMS,
    mlp_hidden_units_factors=MLP_HIDDEN_UNITS_FACTORS,
    dropout_rate=DROPOUT_RATE,
)

print("Total model weights:", tabtransformer_model.count_params())
keras.utils.plot_model(tabtransformer_model, show_shapes=True, rankdir="LR")



history = run_experiment(
    model=tabtransformer_model,
    train_data_file=train_data_file,
    test_data_file=test_data_file,
    num_epochs=NUM_EPOCHS,
    learning_rate=LEARNING_RATE,
    weight_decay=WEIGHT_DECAY,
    batch_size=BATCH_SIZE,
)