geogenie.models package

Submodules

geogenie.models.models module

class geogenie.models.models.MLPRegressor(input_size, width=256, nlayers=10, dropout_prop=0.25, device='cpu', output_width=2, dtype=torch.float32, batch_size=32)[source]

Bases: Module

Define PyTorch MLP Model.

This class defines a PyTorch Multi-Layer Perceptron (MLP) model for regression tasks.

input_size

The number of input features.

Type:

int

width

The width of the hidden layers.

Type:

int

nlayers

The number of hidden layers.

Type:

int

dropout_prop

The dropout proportion.

Type:

float

device

The device to run the model on.

Type:

str

dtype

The data type of the model.

Type:

torch.dtype

batch_size

The batch size for the model.

Type:

int

seqmodel

The sequential model

Type:

torch.nn.Sequential

_define_model(input_size, width, nlayers, dropout_prop, output_width)[source]

Method to define the neural network model.

Parameters:

  • input_size (int) – The number of input features.

  • width (int) – The width of the hidden layers.

  • nlayers (int) – The number of hidden layers.

  • dropout_prop (float) – The dropout proportion.

  • output_width (int) – The number of output features.

Returns:

The sequential model.

Return type:

torch.nn.Sequential

forward(x)[source]

Forward pass through the neural network.

Parameters:

x (torch.Tensor) – The input tensor to the neural network.

Returns:

The output tensor after passing through the network.

Return type:

torch.Tensor

Module contents