TabR
A deep learning model that integrates a KNN component to enhance tabular data predictions through an efficient attention-like mechanism.
Functions
def _initialize_embeddings(weight: Tensor, d: Optional[int]) -> None
Initializes embedding weights with uniform distribution.
Parameters:
weight (Tensor) - Weight tensor to initialize.
d (Optional[int]) - Embedding dimension.
def make_trainable_vector(d: int) -> Parameter
Creates a trainable parameter vector with proper initialization.
Parameters:
d (int) - Vector dimension.
Returns:
Parameter - Trainable parameter vector.
class CLSEmbedding(nn.Module)
Adds a CLS token to the beginning of input sequences.
Parameters:
d_embedding (int) - Embedding dimension.
Input Shape:
(*, seq_len, d_embedding)
Output Shape:
(*, seq_len + 1, d_embedding)
class ResNet(nn.Module)
Residual network with customizable blocks and normalization.
Parameters:
d_in (Optional[int]) - Input dimension.
d_out (Optional[int]) - Output dimension.
n_blocks (int) - Number of residual blocks.
d_block (int) - Block dimension.
dropout (float) - Dropout rate.
d_hidden_multiplier (float) - Hidden dimension multiplier.
n_linear_layers_per_block (int, optional, Default is 2) - Number of linear layers per block.
activation (str, optional, Default is ‘ReLU’) - Activation function.
normalization (str) - Normalization type.
first_normalization (bool) - Whether to apply normalization first.
class LinearEmbeddings(nn.Module)
Linear embeddings for continuous features.
Parameters:
n_features (int) - Number of features.
d_embedding (int) - Embedding dimension.
bias (bool, optional, Default is True) - Whether to use bias.
class PeriodicEmbeddings(nn.Module)
Periodic embeddings using frequency-based encoding.
Parameters:
n_features (int) - Number of features.
n_frequencies (int) - Number of frequencies.
frequency_scale (float) - Frequency scaling factor.
class NLinear(nn.Module)
Feature-wise linear layer.
Parameters:
n_features (int) - Number of features.
d_in (int) - Input dimension.
d_out (int) - Output dimension.
bias (bool, optional, Default is True) - Whether to use bias.
class LREmbeddings(nn.Sequential)
Linear + ReLU embeddings.
Parameters:
n_features (int) - Number of features.
d_embedding (int) - Embedding dimension.
class PLREmbeddings(nn.Sequential)
Periodic + Linear + ReLU embeddings.
Parameters:
n_features (int) - Number of features.
n_frequencies (int) - Number of frequencies.
frequency_scale (float) - Frequency scaling factor.
d_embedding (int) - Embedding dimension.
lite (bool) - Whether to use lite version.
class PBLDEmbeddings(nn.Module)
Periodic + BatchNorm + Linear + Dropout embeddings.
Parameters:
n_features (int) - Number of features.
n_frequencies (int) - Number of frequencies.
frequency_scale (float) - Frequency scaling factor.
d_embedding (int) - Embedding dimension.
lite (bool) - Whether to use lite version.
plr_act_name (str, optional, Default is ‘relu’) - Activation function name.
plr_use_densenet (bool, optional, Default is True) - Whether to use dense connections.
class MLP(nn.Module)
Multi-layer perceptron with SELU activation.
Parameters:
d_in (Optional[int]) - Input dimension.
d_out (Optional[int]) - Output dimension.
n_blocks (int) - Number of blocks.
d_block (int) - Block dimension.
dropout (float) - Dropout rate.
activation (str, optional, Default is ‘SELU’) - Activation function.
def make_module(spec, *args, **kwargs) -> nn.Module
Creates a module based on specification.
Parameters:
spec - Module specification.
args - Positional arguments.
kwargs - Keyword arguments.
Returns:
nn.Module - Created module.
def make_module1(type: str, *args, **kwargs) -> nn.Module
Creates a module of specified type.
Parameters:
type (str) - Module type.
args - Positional arguments.
kwargs - Keyword arguments.
Returns:
nn.Module - Created module.
References:
Yury Gorishniy, Ivan Rubachev, Nikolay Kartashev, Daniil Shlenskii, Akim Kotelnikov, and Artem Babenko. TabR: Tabular Deep Learning Meets Nearest Neighbors in 2023. arXiv:2307.14338 [cs.LG], 2023. https://arxiv.org/abs/2307.14338