Method __init__

exir/tests/models.py:72–86

(self, n_layer: int = 1, output_size: int = 1)

class MLP(nn.Module):
    def __init__(self, n_layer: int = 1, output_size: int = 1) -> None:
        super().__init__()
        self.n_layer = n_layer
        self.output_size = output_size
        # input shape [batch_size, n_layer+output_size]
        # each linear layer reduce the activation dim 1 size by 1.
        self.mlp = torch.nn.Sequential(
            *itertools.chain(
                *(
                    [nn.Linear(i + output_size, i - 1 + output_size)]
                    + ([nn.ReLU()] if i != 1 else [])
                    for i in range(n_layer, 0, -1)
                )
            )
        )

    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
        return self.mlp(inputs)