dataeval.utils.torch.trainer.AETrainer

class dataeval.utils.torch.trainer.AETrainer(model, device='auto', batch_size=8)

A class to train and evaluate an autoencoder<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)

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

Note

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

eval(dataset)

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, epochs=25)

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]

Note

To replace this function with a custom function, do:

AETrainer.train = custom_function