diffusion_models.utils.schemas

Module Contents

class TrainingConfiguration

training_name: str

batch_size: int

learning_rate: float

number_of_epochs: int

log_rate: int = 10

image_rate: int = 100

checkpoint_rate: int = 100

mixed_precision_training: bool = False

gradient_clip: float | None = None

class LogConfiguration

log_rate: int = 10

image_rate: int = 50

number_of_images: int = 5

class BetaSchedulerConfiguration

steps: int

betas: torch.Tensor

alpha_bars: torch.Tensor

class Checkpoint

epoch: int

model_state_dict: Dict[str, Any]

optimizer_state_dict: Dict[str, Any]

scaler: torch.cuda.amp.GradScaler | None

beta_scheduler_config: BetaSchedulerConfiguration

tensorboard_run_name: str | None = None

image_channels: int = 3

loss: float | None = None

classmethod from_file(file_path)

to_file(file_path)

class OldCheckpoint

epoch: int

model_state_dict: Dict[str, Any]

optimizer_state_dict: Dict[str, Any]

scaler: torch.cuda.amp.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)