torch_frame.nn.encoder.LinearModelEncoder

class LinearModelEncoder(out_channels: Optional[int] = None, stats_list: Optional[list[dict[torch_frame.data.stats.StatType, Any]]] = None, stype: Optional[stype] = None, post_module: Optional[Module] = None, na_strategy: Optional[NAStrategy] = None, col_to_model_cfg: Optional[dict[str, torch_frame.config.model.ModelConfig]] = None)

Bases: StypeEncoder

Linear function based encoder with a specified model outputs embedding features. It applies a linear layer torch.nn.Linear(in_channels, out_channels) on each embedding feature (in_channels is the dimensionality of the embedding) and concatenates the output embeddings. The model will also be trained together with the linear layer. Note that the implementation does this for all columns in a batched manner.

Parameters:

col_to_model_cfg (dict) – A dictionary mapping column names to ModelConfig, which specifies a model to map a single-column TensorData object of shape [batch_size, 1, *] into row embeddings of shape [batch_size, 1, model_out_channels].

reset_parameters() → None

Initialize the parameters of post_module.

encode_forward(feat: torch.Tensor | torch_frame.data.multi_nested_tensor.MultiNestedTensor | torch_frame.data.multi_embedding_tensor.MultiEmbeddingTensor | dict[str, torch_frame.data.multi_nested_tensor.MultiNestedTensor], col_names: Optional[list[str]] = None) → Tensor

The main forward function. Maps input feat from TensorFrame (shape [batch_size, num_cols]) into output x of shape [batch_size, num_cols, out_channels].