clarena.backbones.mlp

 1"""
 2The submodule in `backbones` for MLP backbone network.
 3"""
 4
 5__all__ = ["MLP"]
 6
 7from typing import Callable
 8
 9from torch import Tensor, nn
10from torchvision.ops import MLP as TorchvisionMLP
11
12from clarena.backbones import CLBackbone
13
14
15class MLP(CLBackbone):
16    """Multi-Layer Perceptron a.k.a. Fully-Connected Network.
17
18    Modified from `torchvision.ops.MLP` in accordance with this framework.
19    """
20
21    def __init__(
22        self,
23        input_dim: int,
24        hidden_dims: list[int],
25        output_dim: int,
26        activation_layer: Callable[..., nn.Module] | None = nn.ReLU,
27        bias: bool = True,
28        dropout: float = 0.0,
29    ):
30        """Initialise the MLP backbone network.
31
32        **Args:**
33        - **input_dim** (`int`): the input dimension. Any data need to be flattened before going in MLP. Note that it is not required in convolutional networks.
34        - **hidden_dims** (`list[int]`): list of hidden layer dimensions. It can be empty list which means single-layer MLP, and it can be as many layers as you want. Note that it doesn't include the last dimension which we take as output dimension.
35        - **output_dim** (`int`): the output dimension which connects to CL output heads.
36        - **activation_layer** (`Callable[..., torch.nn.Module]` | `None`): activation function of each layer (if not `None`), if `None` this layer won't be used.
37        - **bias** (`bool`): whether to use bias in the linear layer. Default `True`.
38        - **dropout** (`float`): The probability for the dropout layer. Default: 0.0.
39        """
40        super().__init__(output_dim=output_dim)
41
42        self.mlp = TorchvisionMLP(
43            in_channels=input_dim,
44            hidden_channels=hidden_dims + [output_dim],
45            activation_layer=activation_layer,
46            bias=bias,
47            dropout=dropout,
48        )
49
50    def forward(self, input: Tensor, task_id: int | None = None) -> Tensor:
51        """The forward pass for data. It is the same for all tasks.
52
53        **Args:**
54        - **input** (`Tensor`): The input tensor from data.
55        - **task_id** (`int`): the task ID where the data are from. It is just a placeholder for API consistence but never used.
56
57        **Returns:**
58        - The output feature tensor to be passed into heads.
59        """
60        input_flat = input.view(input.size(0), -1)  # flatten before going through MLP
61        feature = self.mlp(input_flat)
62
63        return feature