softmax_bottleneck_model.py

from typing import TYPE_CHECKING

import torch
from torch import nn

from src.config import WEIGHT_INIT_METHODS, WeightInit
from src.utils import Activation

if TYPE_CHECKING:
    from collections.abc import Callable


class SoftmaxBottleneckModel(nn.Module):
    def __init__(
        self,
        output_size: int,
        hidden_size: int,
        num_hidden_layers: int,
        activation: Activation | None = None,
        weight_init: WeightInit | None = None,
    ) -> None:
        super().__init__()
        self.hidden_dim: int = hidden_size
        self.output_dim: int = output_size
        self.fc_in = nn.Linear(output_size, hidden_size)
        self.hidden_layers: list[nn.Linear] = []
        if activation is None:
            activation = Activation(activation_fun=nn.ReLU())
        self.activation = (
            activation.activation_fun
            if activation.dim is None
            else activation.activation_fun(dim=activation.dim)
        )

        for _ in range(num_hidden_layers):
            self.hidden_layers.append(nn.Linear(hidden_size, hidden_size))

        self.fc_out = nn.Linear(hidden_size, output_size, bias=False)

        if weight_init:
            if weight_init not in WEIGHT_INIT_METHODS:
                exc_msg = f"Invalid weight initialization method: {weight_init}. Choose from {list(WEIGHT_INIT_METHODS.keys())}."
                raise ValueError(exc_msg)

            init_fn: Callable[[torch.Tensor], None] = WEIGHT_INIT_METHODS[weight_init]
            init_fn(self.fc_in.weight)
            init_fn(self.fc_out.weight)
            for i in range(len(self.hidden_layers)):
                init_fn(self.hidden_layers[i].weight)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.fc_in(x)
        x = self.activation(x)
        for layer in self.hidden_layers:
            x = layer(x)
            x = self.activation(x)
        return self.fc_out(x)  # type: ignore[no-any-return]