/ludwig/modules/tabnet_modules.py

import torch
import torch.nn as nn

from ludwig.modules.normalization_modules import GhostBatchNormalization
from ludwig.utils.entmax import Entmax15, EntmaxBisect, Sparsemax
from ludwig.utils.torch_utils import LudwigModule


class TabNet(LudwigModule):
    def __init__(
        self,
        input_size: int,
        size: int,
        output_size: int,
        num_steps: int = 1,
        num_total_blocks: int = 4,
        num_shared_blocks: int = 2,
        relaxation_factor: float = 1.5,
        bn_momentum: float = 0.3,
        bn_epsilon: float = 1e-3,
        bn_virtual_bs: int | None = None,
        sparsity: float = 1e-5,
        entmax_mode: str = "sparsemax",
        entmax_alpha: float = 1.5,
    ):
        """TabNet Will output a vector of size output_dim.

        Args:
            input_size: concatenated size of input feature encoder outputs
            size: Embedding feature dimension
            output_size: Output dimension for TabNet
            num_steps: Total number of steps.
            num_total_blocks: Total number of feature transformer blocks.
            num_shared_blocks: Number of shared feature transformer blocks.
            relaxation_factor: >1 will allow features to be used more than once.
            bn_momentum: Batch normalization, momentum.
            bn_epsilon: Batch normalization, epsilon.
            bn_virtual_bs: Virtual batch ize for ghost batch norm.
            entmax_mode: Entmax is a sparse family of probability mapping which generalizes softmax and sparsemax.
                         entmax_mode controls the sparsity.  One of {"sparsemax", "entmax15", "constant", "adaptive"}.
            entmax_alpha: Must be a number between 1.0 and 2.0.  If entmax_mode is "adaptive", entmax_alpha is used
                          as the initial value for the learnable parameter.
        """
        super().__init__()
        self.input_size = input_size
        self.size = size
        self.output_size = output_size
        self.num_steps = num_steps
        self.bn_virtual_bs = bn_virtual_bs
        self.relaxation_factor = relaxation_factor
        self.sparsity = torch.tensor(sparsity)
        self.batch_norm = nn.BatchNorm1d(input_size, momentum=bn_momentum, eps=bn_epsilon)

        kargs = {
            "num_total_blocks": num_total_blocks,
            "num_shared_blocks": num_shared_blocks,
            "bn_momentum": bn_momentum,
            "bn_epsilon": bn_epsilon,
            "bn_virtual_bs": bn_virtual_bs,
        }

        # first feature transformer block is built first
        # to get the shared blocks
        self.feature_transforms = nn.ModuleList([FeatureTransformer(input_size, size + output_size, **kargs)])
        self.attentive_transforms = nn.ModuleList([None])
        for _i in range(num_steps):
            self.feature_transforms.append(
                FeatureTransformer(
                    input_size,
                    size + output_size,
                    **kargs,
                    shared_fc_layers=self.feature_transforms[0].shared_fc_layers,
                )
            )
            # attentive transformers are initialized in build
            # because their outputs size depends on the number
            # of features that we determine by looking at the
            # last dimension of the input tensor
            self.attentive_transforms.append(
                AttentiveTransformer(
                    size, input_size, bn_momentum, bn_epsilon, bn_virtual_bs, entmax_mode, entmax_alpha
                )
            )
        self.final_projection = nn.Linear(output_size, output_size)

        # Register tensors to be used in forward pass. This is needed in order to move these tensors
        # to the correct device (GPU/CPU) during the forward pass.
        self.register_buffer("out_accumulator", torch.zeros(output_size))
        self.register_buffer("aggregated_mask", torch.zeros(input_size))
        self.register_buffer("prior_scales", torch.ones(input_size))

    def forward(self, features: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, list[torch.Tensor]]:
        if features.dim() != 2:
            raise ValueError(f"Expecting incoming tensor to be dim 2, instead dim={features.dim()}")

        # shape notation
        # i_s: input_size
        # s: size
        # o_s: output_size
        # b_s: batch_size
        batch_size = features.shape[0]  # b_s
        # Tile out_accumulator, aggregated_mask, and prior_scales to add batch dimension.
        out_accumulator = torch.tile(self.out_accumulator, (batch_size, 1))
        aggregated_mask = torch.tile(self.aggregated_mask, (batch_size, 1))
        prior_scales = torch.tile(self.prior_scales, (batch_size, 1))
        masks = []
        total_entropy = 0.0

        if batch_size != 1 or not self.training:
            # Skip batch normalization training if the batch size is 1.
            features = self.batch_norm(features)  # [b_s, i_s]
        elif batch_size == 1:
            # We temporarily set the batch_norm module to eval mode as we can't compute the running statistics
            # when the batch size is 1.
            self.batch_norm.eval()
            features = self.batch_norm(features)  # [b_s, i_s]
            self.batch_norm.train()
        masked_features = features

        x = self.feature_transforms[0](masked_features)  # [b_s, s + o_s]

        for step_i in range(1, self.num_steps + 1):
            #########################
            # Attentive Transformer #
            #########################
            # x in following is shape [b_s, s]
            mask_values = self.attentive_transforms[step_i](x[:, self.output_size :], prior_scales)  # [b_s, i_s]

            # relaxation factor 1 forces the feature to be only used once
            prior_scales = prior_scales * (self.relaxation_factor - mask_values)  # [b_s, i_s]

            # entropy is used to penalize the amount of sparsity
            # in feature selection
            if self.sparsity.item() != 0.0:
                total_entropy += (
                    torch.mean(torch.sum(-mask_values * torch.log(mask_values + 0.00001), dim=1)) / self.num_steps
                )

            masks.append(torch.unsqueeze(torch.unsqueeze(mask_values, 0), 3))  # [1, b_s, i_s, 1]

            #######################
            # Feature Transformer #
            #######################
            masked_features = torch.multiply(mask_values, features)

            x = self.feature_transforms[step_i](masked_features)  # [b_s, s + o_s]

            # x in following is shape [b_s, o_s]
            out = nn.functional.relu(x[:, : self.output_size])  # [b_s, o_s]
            out_accumulator += out

            # Aggregated masks are used for visualization of the
            # feature importance attributes.
            scale = torch.sum(out, dim=1, keepdim=True) / self.num_steps
            aggregated_mask += mask_values * scale  # [b_s, i_s]

        final_output = self.final_projection(out_accumulator)  # [b_s, o_s]

        sparsity_loss = torch.multiply(self.sparsity, total_entropy)
        self.update_loss("sparsity_loss", sparsity_loss)

        return final_output, aggregated_mask, masks

    @property
    def input_shape(self) -> torch.Size:
        return torch.Size([self.input_size])

    @property
    def output_shape(self) -> torch.Size:
        return torch.Size([self.output_size])


class FeatureBlock(LudwigModule):
    def __init__(
        self,
        input_size: int,
        size: int,
        apply_glu: bool = True,
        bn_momentum: float = 0.1,
        bn_epsilon: float = 1e-3,
        bn_virtual_bs: int | None = None,
        shared_fc_layer: LudwigModule = None,
    ):
        super().__init__()
        self.input_size = input_size
        self.apply_glu = apply_glu
        self.size = size
        units = size * 2 if apply_glu else size

        # Initialize fc_layer before assigning to shared layer for torchscript compatibilty
        self.fc_layer = nn.Linear(input_size, units, bias=False)
        if shared_fc_layer is not None:
            if shared_fc_layer.weight.shape != self.fc_layer.weight.shape:
                raise RuntimeError(
                    f"shared_fc_layer weight shape {tuple(shared_fc_layer.weight.shape)} doesn't match "
                    f"expected shape {tuple(self.fc_layer.weight.shape)} for input_size={input_size}, size={size}."
                )
            self.fc_layer = shared_fc_layer

        self.batch_norm = GhostBatchNormalization(
            units, virtual_batch_size=bn_virtual_bs, momentum=bn_momentum, epsilon=bn_epsilon
        )

    def forward(self, inputs):
        # shape notation
        # i_s: input_size
        # s: size
        # u: units
        # b_s: batch_size

        # inputs shape [b_s, i_s]
        hidden = self.fc_layer(inputs)  # [b_s, u]
        hidden = self.batch_norm(hidden)  # [b_s, u]
        if self.apply_glu:
            hidden = nn.functional.glu(hidden, dim=-1)  # [bs, s]
        return hidden  # [b_s, 2*s] if apply_glu else [b_s, s]

    @property
    def input_shape(self) -> torch.Size:
        return torch.Size([self.input_size])


class AttentiveTransformer(LudwigModule):
    def __init__(
        self,
        input_size: int,
        size: int,
        bn_momentum: float = 0.1,
        bn_epsilon: float = 1e-3,
        bn_virtual_bs: int | None = None,
        entmax_mode: str = "sparsemax",
        entmax_alpha: float = 1.5,
    ):
        super().__init__()
        self.input_size = input_size
        self.size = size
        self.entmax_mode = entmax_mode
        if entmax_mode == "adaptive":
            self.register_buffer("trainable_alpha", torch.tensor(entmax_alpha, requires_grad=True))
        else:
            self.trainable_alpha = entmax_alpha

        if self.entmax_mode == "sparsemax":
            self.entmax_module = Sparsemax()
        elif self.entmax_mode == "entmax15":
            self.entmax_module = Entmax15()
        else:
            self.entmax_module = EntmaxBisect(alpha=self.trainable_alpha)

        self.feature_block = FeatureBlock(
            input_size,
            size,
            bn_momentum=bn_momentum,
            bn_epsilon=bn_epsilon,
            bn_virtual_bs=bn_virtual_bs,
            apply_glu=False,
        )

    def forward(self, inputs, prior_scales):
        # shape notation
        # i_s: input_size
        # s: size
        # b_s: batch_size

        # inputs shape [b_s, i_s], prior_scales shape [b_s, s]
        hidden = self.feature_block(inputs)  # [b_s, s]
        hidden = hidden * prior_scales  # [b_s, s]

        # removing the mean to try to avoid numerical instability
        # https://github.com/tensorflow/addons/issues/2314
        # https://github.com/tensorflow/tensorflow/pull/21183/files
        # In (Arik and Pfister, 2019), they call the logits z.
        # The mean(logits) can be substracted from logits to make the algorithm
        # more numerically stable. the instability in this algorithm comes mostly
        # from the z_cumsum. Substacting the mean will cause z_cumsum to be close
        # to zero.
        # hidden = hidden - tf.math.reduce_mean(hidden, axis=1)[:, tf.newaxis]
        return self.entmax_module(hidden)

    @property
    def input_shape(self) -> torch.Size:
        return torch.Size([self.input_size])

    @property
    def output_shape(self) -> torch.Size:
        return torch.Size([self.size])


# adapted and modified from:
# https://github.com/ostamand/tensorflow-tabnet/blob/master/tabnet/models/transformers.py
class FeatureTransformer(LudwigModule):
    def __init__(
        self,
        input_size: int,
        size: int,
        shared_fc_layers: list | None = None,
        num_total_blocks: int = 4,
        num_shared_blocks: int = 2,
        bn_momentum: float = 0.1,
        bn_epsilon: float = 1e-3,
        bn_virtual_bs: int | None = None,
    ):
        super().__init__()
        if shared_fc_layers is None:
            shared_fc_layers = []
        self.input_size = input_size
        self.num_total_blocks = num_total_blocks
        self.num_shared_blocks = num_shared_blocks
        self.size = size

        kwargs = {
            "bn_momentum": bn_momentum,
            "bn_epsilon": bn_epsilon,
            "bn_virtual_bs": bn_virtual_bs,
        }

        # build blocks
        self.blocks = nn.ModuleList()
        for n in range(num_total_blocks):
            # Ensure the sizes fed into FeatureBlock are correct regardless of presence of shared_fc_layer
            if n == 0:
                in_features = input_size
            else:
                in_features = size

            if shared_fc_layers and n < len(shared_fc_layers):
                self.blocks.append(FeatureBlock(in_features, size, **kwargs, shared_fc_layer=shared_fc_layers[n]))
            else:
                self.blocks.append(FeatureBlock(in_features, size, **kwargs))

    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
        # shape notation
        # i_s: input_size
        # s: size
        # b_s: batch_size

        # inputs shape [b_s, i_s]
        hidden = self.blocks[0](inputs)  # [b_s, s]
        for n in range(1, self.num_total_blocks):
            hidden = (self.blocks[n](hidden) + hidden) * (0.5**0.5)  # [b_s, s]
        return hidden  # [b_s, s]

    @property
    def shared_fc_layers(self):
        return [self.blocks[i].fc_layer for i in range(self.num_shared_blocks)]

    @property
    def input_shape(self) -> torch.Size:
        return torch.Size([self.input_size])

    @property
    def output_shape(self) -> torch.Size:
        return torch.Size([self.size])