Simple dense net

Simple dense neural network.

SimpleDenseNet

Bases: Module

A simple fully-connected neural net for computing predictions.

Source code in src/models/components/simple_dense_net.py

class SimpleDenseNet(nn.Module):
    """A simple fully-connected neural net for computing predictions."""

    def __init__(
        self,
        input_size: int = 784,
        lin1_size: int = 256,
        lin2_size: int = 256,
        lin3_size: int = 256,
        output_size: int = 10,
    ) -> None:
        """Initialize a `SimpleDenseNet` module.

        Args:
            input_size: The number of input features.
            lin1_size: The number of output features of the first linear layer.
            lin2_size: The number of output features of the second linear layer.
            lin3_size: The number of output features of the third linear layer.
            output_size: The number of output features of the final linear layer.
        """
        super().__init__()

        self.model = nn.Sequential(
            nn.Linear(input_size, lin1_size),
            nn.BatchNorm1d(lin1_size),
            nn.ReLU(),
            nn.Linear(lin1_size, lin2_size),
            nn.BatchNorm1d(lin2_size),
            nn.ReLU(),
            nn.Linear(lin2_size, lin3_size),
            nn.BatchNorm1d(lin3_size),
            nn.ReLU(),
            nn.Linear(lin3_size, output_size),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Perform a single forward pass through the network.

        Args:
            x: The input tensor.

        Returns:
            A tensor of predictions.
        """
        batch_size, channels, width, height = x.size()

        # (batch, 1, width, height) -> (batch, 1*width*height)
        x = x.view(batch_size, -1)

        return self.model(x)

__init__(input_size=784, lin1_size=256, lin2_size=256, lin3_size=256, output_size=10)

Initialize a SimpleDenseNet module.

Parameters:

Name	Type	Description	Default
input_size	int	The number of input features.	784
lin1_size	int	The number of output features of the first linear layer.	256
lin2_size	int	The number of output features of the second linear layer.	256
lin3_size	int	The number of output features of the third linear layer.	256
output_size	int	The number of output features of the final linear layer.	10

Source code in src/models/components/simple_dense_net.py

def __init__(
    self,
    input_size: int = 784,
    lin1_size: int = 256,
    lin2_size: int = 256,
    lin3_size: int = 256,
    output_size: int = 10,
) -> None:
    """Initialize a `SimpleDenseNet` module.

    Args:
        input_size: The number of input features.
        lin1_size: The number of output features of the first linear layer.
        lin2_size: The number of output features of the second linear layer.
        lin3_size: The number of output features of the third linear layer.
        output_size: The number of output features of the final linear layer.
    """
    super().__init__()

    self.model = nn.Sequential(
        nn.Linear(input_size, lin1_size),
        nn.BatchNorm1d(lin1_size),
        nn.ReLU(),
        nn.Linear(lin1_size, lin2_size),
        nn.BatchNorm1d(lin2_size),
        nn.ReLU(),
        nn.Linear(lin2_size, lin3_size),
        nn.BatchNorm1d(lin3_size),
        nn.ReLU(),
        nn.Linear(lin3_size, output_size),
    )

forward(x)

Perform a single forward pass through the network.

Parameters:

Name	Type	Description	Default
x	Tensor	The input tensor.	required

Returns:

Type	Description
Tensor	A tensor of predictions.

Source code in src/models/components/simple_dense_net.py

def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Perform a single forward pass through the network.

    Args:
        x: The input tensor.

    Returns:
        A tensor of predictions.
    """
    batch_size, channels, width, height = x.size()

    # (batch, 1, width, height) -> (batch, 1*width*height)
    x = x.view(batch_size, -1)

    return self.model(x)