torch_frame.nn.models.FTTransformer

class FTTransformer(channels: int, out_channels: int, num_layers: int, col_stats: dict[str, dict[StatType, Any]], col_names_dict: dict[torch_frame.stype, list[str]], stype_encoder_dict: dict[torch_frame.stype, StypeEncoder] | None = None)

Bases: Module

The FT-Transformer model introduced in the “Revisiting Deep Learning Models for Tabular Data” paper.

Note

For an example of using FTTransformer, see examples/revisiting.py.

Parameters:

• channels (int) – Hidden channel dimensionality

• out_channels (int) – Output channels dimensionality

• num_layers (int) – Number of layers. (default: 3)

• col_stats (dict[str,dict[torch_frame.data.stats.StatType,Any]]) – A dictionary that maps column name into stats. Available as dataset.col_stats.

• col_names_dict (dict[torch_frame.stype, list[str]]) – A dictionary that maps stype to a list of column names. The column names are sorted based on the ordering that appear in tensor_frame.feat_dict. Available as tensor_frame.col_names_dict.

• stype_encoder_dict – (dict[torch_frame.stype, torch_frame.nn.encoder.StypeEncoder], optional): A dictionary mapping stypes into their stype encoders. (default: None, will call torch_frame.nn.encoder.EmbeddingEncoder() for categorical feature and torch_frame.nn.encoder.LinearEncoder() for numerical feature)

forward(tf: TensorFrame) → Tensor

Transforming TensorFrame object into output prediction.

Parameters:

tf (TensorFrame) – Input TensorFrame object.

Returns:

Output of shape [batch_size, out_channels].

Return type:

torch.Tensor