torch_frame.nn.models.MLP

class MLP(channels: int, out_channels: int, num_layers: int, col_stats: dict[str, dict[torch_frame.data.stats.StatType, Any]], col_names_dict: dict[torch_frame._stype.stype, list[str]], stype_encoder_dict: Optional[dict[torch_frame._stype.stype, torch_frame.nn.encoder.stype_encoder.StypeEncoder]] = None, normalization: str | None = 'layer_norm', dropout_prob: float = 0.2)[source]

The light-weight MLP model that mean-pools column embeddings and applies MLP over it.

Parameters:

channels (int) – The number of channels in the backbone layers.
out_channels (int) – The number of output channels in the decoder.
num_layers (int) – The number of layers in the backbone.
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 EmbeddingEncoder() for categorical feature and LinearEncoder() for numerical feature)
normalization (str, optional) – The type of normalization to use. batch_norm, layer_norm, or None. (default: layer_norm)
dropout_prob (float) – The dropout probability (default: 0.2).

Transforming TensorFrame object into output prediction.