oumi.models

oumi.models#

Models module for the Oumi (Open Universal Machine Intelligence) library.

This module provides various model implementations for use in the Oumi framework. These models are designed for different machine learning tasks and can be used with the datasets and training pipelines provided by Oumi.

Available models:
  • MLPEncoder: A Multi-Layer Perceptron (MLP)

    encoder model.

  • CNNClassifier: A simple ConvNet for

    image classification e.g., can be used for MNIST digits classification.

Each model is implemented as a separate class, inheriting from appropriate base classes in the Oumi framework.

Example

>>> from oumi.models import MLPEncoder
>>> encoder = MLPEncoder(input_dim=784, hidden_dim=256, output_dim=10)
>>> output = encoder(input_data)

Note

For detailed information on each model, please refer to their respective class documentation.

class oumi.models.CNNClassifier(image_width: int, image_height: int, *, in_channels: int = 3, output_dim: int = 10, kernel_size: int = 5)[source]#

Bases: BaseModel

A simple ConvNet for classification of small fixed-size images.

property criterion: Callable#

Returns the criterion function to compute loss.

forward(images: Tensor, labels: LongTensor | None = None, **kwargs) → dict[str, Tensor][source]#

Forward pass of the model.

class oumi.models.MLPEncoder(input_dim: int = 768, hidden_dim: int = 128, output_dim: int = 10)[source]#

Bases: BaseModel

property criterion: Callable#

Returns the criterion function for the MLP model.

The criterion function is used to compute the loss during training.

Returns:

The cross-entropy loss function.

Return type:

torch.nn.CrossEntropyLoss

forward(input_ids: LongTensor, labels: LongTensor | None = None, **kwargs) → dict[str, Tensor][source]#

Forward pass of the MLP model.

Parameters:

  • input_ids (torch.LongTensor) – The input tensor of shape (batch_size, sequence_length).

  • labels (torch.LongTensor, optional) – The target labels tensor of shape (batch_size,).

  • **kwargs – Additional keyword arguments provided by the tokenizer. Not used in this model.

Returns:

A dictionary containing the model outputs.

The dictionary has the following keys:

  • ”logits” (torch.Tensor): The output logits tensor of shape (batch_size, num_classes).

  • ”loss” (torch.Tensor, optional): The computed loss tensor if labels is not None.

Return type:

dict

Subpackages#