diffusion_models.utils.schemas

Module Contents

class TrainingConfiguration

A training configuration for simple experiment management.

training_name: str

The name of the training.

batch_size: int

The batch size used for training.

learning_rate: float

The learning rate of the training.

number_of_epochs: int

The number of epoch used for training.

checkpoint_rate: int = 100

The rate at which checkpoints are saved..

mixed_precision_training: bool = False

Whether or not to use automatic mixed precision training.

gradient_clip: float | None = None

Whether or not to clip gradients.

class LogConfiguration

An object to manage logging configuration.

log_rate: int = 10

The rate at which training metrics are logged.

image_rate: int = 50

The rate at which images are generated for visualization. This can be used to validate model performance.

number_of_images: int = 5

The number of images to generate.

class BetaSchedulerConfiguration

A simplified beta scheduler configuration.

steps: int

The number of steps in the beta scheduler.

betas: Tensor

The beta values.

alpha_bars: Tensor

The alpha bar values.

class Checkpoint

A simplified checkpoint framework for easy saving and loading.

epoch: int

The current epoch.

model_state_dict: Dict[str, Any]

The model state dict.

optimizer_state_dict: Dict[str, Any]

The optimizer state dict.

beta_scheduler_config: BetaSchedulerConfiguration

The beta scheduler configuration.

scaler: GradScaler | None = None

The GradScaler instance.

tensorboard_run_name: str | None = None

The name of the tensorboard run.

image_channels: int = 3

The number of image channels used in the training.

loss: float | None = None

The final loss value recorded.

Note

This is a legacy parameter and will be removed in a future release.

classmethod from_file(file_path, map_location=None)

Load and instantiate a checkpoint from a file.

Parameters:

• file_path (str) – The path to the checkpoint file.

• map_location (Optional[str]) – A function, torch. device, string or a dict specifying how to remap storage location.

Returns:

A checkpoint instance.

Return type:

Checkpoint

to_file(file_path)

Saves a checkpoint to a file.

class OldCheckpoint

epoch: int

model_state_dict: Dict[str, Any]

optimizer_state_dict: Dict[str, Any]

scaler: GradScaler | None

tensorboard_run_name: str | None = None

loss: float | None = None

classmethod from_file(file_path)

to_file(file_path)

to_new_checkpoint(beta_scheduler)

class Timestep

current: Tensor

previous: Tensor | None = None