PyTorch Backend#

Models#

class dataeval.torch.models.AriaAutoencoder(channels=3)#

An autoencoder model with a separate encoder and decoder.

Parameters:

channels (int, default 3) – Number of input channels

class dataeval.torch.models.Decoder(channels)#

A simple decoder to be used in an autoencoder model.

This is the decoder used by the AriaAutoencoder model.

Parameters:

channels (int) – Number of output channels

class dataeval.torch.models.Encoder(channels=3)#

A simple encoder to be used in an autoencoder model.

This is the encoder used by the AriaAutoencoder model.

Parameters:

channels (int, default 3) – Number of input channels

Trainer#

class dataeval.torch.trainer.AETrainer(model: Module, device: str | device = 'auto', batch_size: int = 8)#

A class to train and evaluate an autoencoder model.

Parameters:

  • model (nn.Module) – The model to be trained.

  • device (str or torch.device, default "auto") – The hardware device to use for training. If “auto”, the device will be set to “cuda” if available, otherwise “cpu”.

  • batch_size (int, default 8) – The number of images to process in a batch.

encode(dataset: Dataset) → Tensor#

Create image embeddings for the dataset using the model’s encoder.

If the model has an encode method, it will be used; otherwise, model.forward will be used.

Parameters:

dataset (Dataset) – The dataset to encode. Torch Dataset containing images in the first return position.

Returns:

Data encoded by the model

Return type:

torch.Tensor

Notes

This function should be run after the model has been trained and evaluated.

eval(dataset: Dataset) → float#

Basic image reconstruction evaluation function for Autoencoder models

Uses torch.nn.MSELoss as default loss function.

Parameters:

dataset (Dataset) – The dataset to evaluate on. Torch Dataset containing images in the first return position.

Returns:

Total reconstruction loss over the entire dataset

Return type:

float

Note

To replace this function with a custom function, do

AETrainer.eval = custom_function

train(dataset: Dataset, epochs: int = 25) → list[float]#

Basic image reconstruction training function for Autoencoder models

Uses torch.optim.Adam and torch.nn.MSELoss as default hyperparameters

Parameters:

  • dataset (Dataset) – The dataset to train on. Torch Dataset containing images in the first return position.

  • epochs (int, default 25) – Number of full training loops

Returns:

A list of average loss values for each epoch.

Return type:

List[float]

Notes

To replace this function with a custom function, do

AETrainer.train = custom_function

On this page