kosmos.ml.models.neural_network

Classes

class NeuralNetwork(input_dim: int, output_dim: int, hidden_layers: list[int])

Bases: kosmos.ml.models.model.Model

Feedforward neural network.

Initialize the neural network.

Parameters:

• input_dim (int) – The input dimension of the model.

• output_dim (int) – The output dimension of the model.

• hidden_layers (list[int]) – Sizes of the hidden layers. The length of the list defines the number of hidden layers, and each element specifies the size of a layer.

Methods

forward(x: torch.Tensor) → torch.Tensor

Perform a forward pass.

Parameters:

x (torch.Tensor) – Input tensor.

Returns:

Output tensor.

Return type:

torch.Tensor