| """Implements a two layer Neural Network.""" |
|
|
| from torch.nn import Module, Linear, ReLU |
| from src.utils.mapper import configmapper |
|
|
|
|
| @configmapper.map("models", "two_layer_nn") |
| class TwoLayerNN(Module): |
| """Implements two layer neural network. |
| |
| Methods: |
| forward(x_input): Returns the output of the neural network. |
| """ |
|
|
| def __init__(self, embedding, dims): |
| """Construct the two layer Neural Network. |
| |
| This method is used to initialize the two layer neural network, |
| with a given embedding type and corresponding arguments. |
| |
| Args: |
| embedding (torch.nn.Module): The embedding layer for the model. |
| dims (list): List of dimensions for the neural network, input to output. |
| """ |
| super(TwoLayerNN, self).__init__() |
|
|
| self.embedding = embedding |
| self.linear1 = Linear(dims[0], dims[1]) |
| self.relu = ReLU() |
| self.linear2 = Linear(dims[1], dims[2]) |
|
|
| def forward(self, x_input): |
| """ |
| Return the output of the neural network for an input. |
| |
| Args: |
| x_input (torch.Tensor): The input tensor to the neural network. |
| |
| Returns: |
| x_output (torch.Tensor): The output tensor for the neural network. |
| """ |
| output = self.embedding(x_input) |
| output = self.linear1(output) |
| output = self.relu(output) |
| x_output = self.linear2(output) |
| return x_output |
|
|
