kosmos.ml.models.neural_network

Classes

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

Bases: kosmos.ml.models.model.Model

Feedforward neural network.

hidden_layers

Sizes of the hidden layers.

Type:

list[int]

model

The sequential model containing all layers.

Type:

nn.Sequential

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[source]

Perform a forward pass.

Parameters:

x (torch.Tensor) – Input tensor.

Returns:

Output tensor.

Return type:

torch.Tensor