gt4sd.training_pipelines.diffusion.core module

Diffusion training utilities. Code adapted from: https://github.com/huggingface/diffusers/blob/main/examples/unconditional_image_generation/train_unconditional.py

Summary

Classes:

DiffusionDataArguments	Data arguments related to diffusion trainer.
DiffusionForVisionTrainingPipeline	Diffusion training pipelines for image generation.
DiffusionModelArguments	Model arguments related to Diffusion trainer.
DiffusionSavingArguments	Saving arguments related to Diffusion trainer.
DiffusionTrainingArguments	Training arguments related to Diffusion trainer.

Reference

class DiffusionForVisionTrainingPipeline

Bases: TrainingPipeline

Diffusion training pipelines for image generation.

train(training_args, model_args, dataset_args)

Generic training function for Diffusion models.

Parameters

• training_args (Dict[str, Any]) – training arguments passed to the configuration.

• model_args (Dict[str, Any]) – model arguments passed to the configuration.

• dataset_args (Dict[str, Any]) – dataset arguments passed to the configuration.

Raises

NotImplementedError – the generic trainer does not implement the pipeline.

Return type

None

__annotations__ = {}

__doc__ = 'Diffusion training pipelines for image generation.'

__module__ = 'gt4sd.training_pipelines.diffusion.core'

class DiffusionDataArguments(dataset_name='', dataset_config_name='', train_data_dir='', resolution=64, train_batch_size=16, eval_batch_size=16, num_epochs=100)

Bases: TrainingPipelineArguments

Data arguments related to diffusion trainer.

__name__ = 'DiffusionDataArguments'

dataset_name: str = ''

dataset_config_name: str = ''

train_data_dir: str = ''

resolution: int = 64

train_batch_size: int = 16

eval_batch_size: int = 16

num_epochs: int = 100

__annotations__ = {'dataset_config_name': <class 'str'>, 'dataset_name': <class 'str'>, 'eval_batch_size': <class 'int'>, 'num_epochs': <class 'int'>, 'resolution': <class 'int'>, 'train_batch_size': <class 'int'>, 'train_data_dir': <class 'str'>}

__dataclass_fields__ = {'dataset_config_name': Field(name='dataset_config_name',type=<class 'str'>,default='',default_factory=<dataclasses._MISSING_TYPE object>,init=True,repr=True,hash=None,compare=True,metadata=mappingproxy({'help': 'Dataset config name.'}),kw_only=False,_field_type=_FIELD), 'dataset_name': Field(name='dataset_name',type=<class 'str'>,default='',default_factory=<dataclasses._MISSING_TYPE object>,init=True,repr=True,hash=None,compare=True,metadata=mappingproxy({'help': 'Dataset name.'}),kw_only=False,_field_type=_FIELD), 'eval_batch_size': Field(name='eval_batch_size',type=<class 'int'>,default=16,default_factory=<dataclasses._MISSING_TYPE object>,init=True,repr=True,hash=None,compare=True,metadata=mappingproxy({'help': 'Eval batch size.'}),kw_only=False,_field_type=_FIELD), 'num_epochs': Field(name='num_epochs',type=<class 'int'>,default=100,default_factory=<dataclasses._MISSING_TYPE object>,init=True,repr=True,hash=None,compare=True,metadata=mappingproxy({'help': 'Number of epochs.'}),kw_only=False,_field_type=_FIELD), 'resolution': Field(name='resolution',type=<class 'int'>,default=64,default_factory=<dataclasses._MISSING_TYPE object>,init=True,repr=True,hash=None,compare=True,metadata=mappingproxy({'help': 'Resolution.'}),kw_only=False,_field_type=_FIELD), 'train_batch_size': Field(name='train_batch_size',type=<class 'int'>,default=16,default_factory=<dataclasses._MISSING_TYPE object>,init=True,repr=True,hash=None,compare=True,metadata=mappingproxy({'help': 'Train batch size.'}),kw_only=False,_field_type=_FIELD), 'train_data_dir': Field(name='train_data_dir',type=<class 'str'>,default='',default_factory=<dataclasses._MISSING_TYPE object>,init=True,repr=True,hash=None,compare=True,metadata=mappingproxy({'help': 'Train data directory.'}),kw_only=False,_field_type=_FIELD)}

__dataclass_params__ = _DataclassParams(init=True,repr=True,eq=True,order=False,unsafe_hash=False,frozen=False)

__doc__ = 'Data arguments related to diffusion trainer.'

__eq__(other)

Return self==value.

__hash__ = None

__init__(dataset_name='', dataset_config_name='', train_data_dir='', resolution=64, train_batch_size=16, eval_batch_size=16, num_epochs=100)

__match_args__ = ('dataset_name', 'dataset_config_name', 'train_data_dir', 'resolution', 'train_batch_size', 'eval_batch_size', 'num_epochs')

__module__ = 'gt4sd.training_pipelines.diffusion.core'

__repr__()

Return repr(self).

class DiffusionModelArguments(model_path, training_name, num_train_timesteps=1000, learning_rate=0.0001, lr_scheduler='cosine', lr_warm_up_steps=500, adam_beta1=0.95, adam_beta2=0.999, adam_weight_decay=1e-06, adam_epsilon=1e-08, gradient_accumulation_steps=1, in_channels=3, out_channels=3, layers_per_block=2)

Bases: TrainingPipelineArguments

Model arguments related to Diffusion trainer.

__name__ = 'DiffusionModelArguments'

model_path: str

training_name: str

num_train_timesteps: int = 1000

learning_rate: float = 0.0001

lr_scheduler: str = 'cosine'

lr_warm_up_steps: int = 500

adam_beta1: float = 0.95

adam_beta2: float = 0.999

adam_weight_decay: float = 1e-06

adam_epsilon: float = 1e-08

gradient_accumulation_steps: int = 1

in_channels: int = 3

out_channels: int = 3

layers_per_block: int = 2

__annotations__ = {'adam_beta1': <class 'float'>, 'adam_beta2': <class 'float'>, 'adam_epsilon': <class 'float'>, 'adam_weight_decay': <class 'float'>, 'gradient_accumulation_steps': <class 'int'>, 'in_channels': <class 'int'>, 'layers_per_block': <class 'int'>, 'learning_rate': <class 'float'>, 'lr_scheduler': <class 'str'>, 'lr_warm_up_steps': <class 'int'>, 'model_path': <class 'str'>, 'num_train_timesteps': <class 'int'>, 'out_channels': <class 'int'>, 'training_name': <class 'str'>}

__dataclass_fields__ = {'adam_beta1': Field(name='adam_beta1',type=<class 'float'>,default=0.95,default_factory=<dataclasses._MISSING_TYPE object>,init=True,repr=True,hash=None,compare=True,metadata=mappingproxy({'help': 'Adam beta1.'}),kw_only=False,_field_type=_FIELD), 'adam_beta2': Field(name='adam_beta2',type=<class 'float'>,default=0.999,default_factory=<dataclasses._MISSING_TYPE object>,init=True,repr=True,hash=None,compare=True,metadata=mappingproxy({'help': 'Adam beta2.'}),kw_only=False,_field_type=_FIELD), 'adam_epsilon': Field(name='adam_epsilon',type=<class 'float'>,default=1e-08,default_factory=<dataclasses._MISSING_TYPE object>,init=True,repr=True,hash=None,compare=True,metadata=mappingproxy({'help': 'Adam eps.'}),kw_only=False,_field_type=_FIELD), 'adam_weight_decay': Field(name='adam_weight_decay',type=<class 'float'>,default=1e-06,default_factory=<dataclasses._MISSING_TYPE object>,init=True,repr=True,hash=None,compare=True,metadata=mappingproxy({'help': 'Adam weights decay.'}),kw_only=False,_field_type=_FIELD), 'gradient_accumulation_steps': Field(name='gradient_accumulation_steps',type=<class 'int'>,default=1,default_factory=<dataclasses._MISSING_TYPE object>,init=True,repr=True,hash=None,compare=True,metadata=mappingproxy({'help': 'Gradient accumulation steps.'}),kw_only=False,_field_type=_FIELD), 'in_channels': Field(name='in_channels',type=<class 'int'>,default=3,default_factory=<dataclasses._MISSING_TYPE object>,init=True,repr=True,hash=None,compare=True,metadata=mappingproxy({'help': 'Input channels.'}),kw_only=False,_field_type=_FIELD), 'layers_per_block': Field(name='layers_per_block',type=<class 'int'>,default=2,default_factory=<dataclasses._MISSING_TYPE object>,init=True,repr=True,hash=None,compare=True,metadata=mappingproxy({'help': 'Layers per block.'}),kw_only=False,_field_type=_FIELD), 'learning_rate': Field(name='learning_rate',type=<class 'float'>,default=0.0001,default_factory=<dataclasses._MISSING_TYPE object>,init=True,repr=True,hash=None,compare=True,metadata=mappingproxy({'help': 'Learning rate.'}),kw_only=False,_field_type=_FIELD), 'lr_scheduler': Field(name='lr_scheduler',type=<class 'str'>,default='cosine',default_factory=<dataclasses._MISSING_TYPE object>,init=True,repr=True,hash=None,compare=True,metadata=mappingproxy({'help': 'Learning rate scheduler.'}),kw_only=False,_field_type=_FIELD), 'lr_warm_up_steps': Field(name='lr_warm_up_steps',type=<class 'int'>,default=500,default_factory=<dataclasses._MISSING_TYPE object>,init=True,repr=True,hash=None,compare=True,metadata=mappingproxy({'help': 'Learning rate warm up steps.'}),kw_only=False,_field_type=_FIELD), 'model_path': Field(name='model_path',type=<class 'str'>,default=<dataclasses._MISSING_TYPE object>,default_factory=<dataclasses._MISSING_TYPE object>,init=True,repr=True,hash=None,compare=True,metadata=mappingproxy({'help': 'Path to the model file.'}),kw_only=False,_field_type=_FIELD), 'num_train_timesteps': Field(name='num_train_timesteps',type=<class 'int'>,default=1000,default_factory=<dataclasses._MISSING_TYPE object>,init=True,repr=True,hash=None,compare=True,metadata=mappingproxy({'help': 'Number of noise steps.'}),kw_only=False,_field_type=_FIELD), 'out_channels': Field(name='out_channels',type=<class 'int'>,default=3,default_factory=<dataclasses._MISSING_TYPE object>,init=True,repr=True,hash=None,compare=True,metadata=mappingproxy({'help': 'Output channels.'}),kw_only=False,_field_type=_FIELD), 'training_name': Field(name='training_name',type=<class 'str'>,default=<dataclasses._MISSING_TYPE object>,default_factory=<dataclasses._MISSING_TYPE object>,init=True,repr=True,hash=None,compare=True,metadata=mappingproxy({'help': 'Name of the training run.'}),kw_only=False,_field_type=_FIELD)}

__dataclass_params__ = _DataclassParams(init=True,repr=True,eq=True,order=False,unsafe_hash=False,frozen=False)

__doc__ = 'Model arguments related to Diffusion trainer.'

__eq__(other)

Return self==value.

__hash__ = None

__init__(model_path, training_name, num_train_timesteps=1000, learning_rate=0.0001, lr_scheduler='cosine', lr_warm_up_steps=500, adam_beta1=0.95, adam_beta2=0.999, adam_weight_decay=1e-06, adam_epsilon=1e-08, gradient_accumulation_steps=1, in_channels=3, out_channels=3, layers_per_block=2)

__match_args__ = ('model_path', 'training_name', 'num_train_timesteps', 'learning_rate', 'lr_scheduler', 'lr_warm_up_steps', 'adam_beta1', 'adam_beta2', 'adam_weight_decay', 'adam_epsilon', 'gradient_accumulation_steps', 'in_channels', 'out_channels', 'layers_per_block')

__module__ = 'gt4sd.training_pipelines.diffusion.core'

__repr__()

Return repr(self).

class DiffusionTrainingArguments(local_rank=-1, output_dir='ddpm-cifar10-32', logging_dir='logs/', overwrite_output_dir=False, cache_dir='.cache/', save_images_epochs=10, save_model_epochs=10, use_ema=True, ema_inv_gamma=1.0, ema_power=0.75, ema_max_decay=0.9999, mixed_precision='no', use_auth_token=False, dummy_training=False, is_sampling=True)

Bases: TrainingPipelineArguments

Training arguments related to Diffusion trainer.

__name__ = 'DiffusionTrainingArguments'

local_rank: int = -1

output_dir: str = 'ddpm-cifar10-32'

logging_dir: str = 'logs/'

overwrite_output_dir: bool = False

cache_dir: str = '.cache/'

save_images_epochs: int = 10

save_model_epochs: int = 10

use_ema: bool = True

ema_inv_gamma: float = 1.0

ema_power: float = 0.75

ema_max_decay: float = 0.9999

mixed_precision: str = 'no'

use_auth_token: bool = False

dummy_training: bool = False

is_sampling: bool = True

__annotations__ = {'cache_dir': <class 'str'>, 'dummy_training': <class 'bool'>, 'ema_inv_gamma': <class 'float'>, 'ema_max_decay': <class 'float'>, 'ema_power': <class 'float'>, 'is_sampling': <class 'bool'>, 'local_rank': <class 'int'>, 'logging_dir': <class 'str'>, 'mixed_precision': <class 'str'>, 'output_dir': <class 'str'>, 'overwrite_output_dir': <class 'bool'>, 'save_images_epochs': <class 'int'>, 'save_model_epochs': <class 'int'>, 'use_auth_token': <class 'bool'>, 'use_ema': <class 'bool'>}

__dataclass_fields__ = {'cache_dir': Field(name='cache_dir',type=<class 'str'>,default='.cache/',default_factory=<dataclasses._MISSING_TYPE object>,init=True,repr=True,hash=None,compare=True,metadata=mappingproxy({'help': 'Cache directory.'}),kw_only=False,_field_type=_FIELD), 'dummy_training': Field(name='dummy_training',type=<class 'bool'>,default=False,default_factory=<dataclasses._MISSING_TYPE object>,init=True,repr=True,hash=None,compare=True,metadata=mappingproxy({'help': 'Run dummy training to test the pipeline.'}),kw_only=False,_field_type=_FIELD), 'ema_inv_gamma': Field(name='ema_inv_gamma',type=<class 'float'>,default=1.0,default_factory=<dataclasses._MISSING_TYPE object>,init=True,repr=True,hash=None,compare=True,metadata=mappingproxy({'help': 'Ema inverse gamma.'}),kw_only=False,_field_type=_FIELD), 'ema_max_decay': Field(name='ema_max_decay',type=<class 'float'>,default=0.9999,default_factory=<dataclasses._MISSING_TYPE object>,init=True,repr=True,hash=None,compare=True,metadata=mappingproxy({'help': 'Ema max delay.'}),kw_only=False,_field_type=_FIELD), 'ema_power': Field(name='ema_power',type=<class 'float'>,default=0.75,default_factory=<dataclasses._MISSING_TYPE object>,init=True,repr=True,hash=None,compare=True,metadata=mappingproxy({'help': 'Ema power.'}),kw_only=False,_field_type=_FIELD), 'is_sampling': Field(name='is_sampling',type=<class 'bool'>,default=True,default_factory=<dataclasses._MISSING_TYPE object>,init=True,repr=True,hash=None,compare=True,metadata=mappingproxy({'help': 'Run sampling.'}),kw_only=False,_field_type=_FIELD), 'local_rank': Field(name='local_rank',type=<class 'int'>,default=-1,default_factory=<dataclasses._MISSING_TYPE object>,init=True,repr=True,hash=None,compare=True,metadata=mappingproxy({'help': 'Local rank of the process.'}),kw_only=False,_field_type=_FIELD), 'logging_dir': Field(name='logging_dir',type=<class 'str'>,default='logs/',default_factory=<dataclasses._MISSING_TYPE object>,init=True,repr=True,hash=None,compare=True,metadata=mappingproxy({'help': 'Logging directory.'}),kw_only=False,_field_type=_FIELD), 'mixed_precision': Field(name='mixed_precision',type=<class 'str'>,default='no',default_factory=<dataclasses._MISSING_TYPE object>,init=True,repr=True,hash=None,compare=True,metadata=mappingproxy({'help': "Mixed precision. Choose from 'no', 'fp16', 'bf16'."}),kw_only=False,_field_type=_FIELD), 'output_dir': Field(name='output_dir',type=<class 'str'>,default='ddpm-cifar10-32',default_factory=<dataclasses._MISSING_TYPE object>,init=True,repr=True,hash=None,compare=True,metadata=mappingproxy({'help': 'Output directory.'}),kw_only=False,_field_type=_FIELD), 'overwrite_output_dir': Field(name='overwrite_output_dir',type=<class 'bool'>,default=False,default_factory=<dataclasses._MISSING_TYPE object>,init=True,repr=True,hash=None,compare=True,metadata=mappingproxy({'help': 'Overwrite output directory.'}),kw_only=False,_field_type=_FIELD), 'save_images_epochs': Field(name='save_images_epochs',type=<class 'int'>,default=10,default_factory=<dataclasses._MISSING_TYPE object>,init=True,repr=True,hash=None,compare=True,metadata=mappingproxy({'help': 'Save images every n epochs.'}),kw_only=False,_field_type=_FIELD), 'save_model_epochs': Field(name='save_model_epochs',type=<class 'int'>,default=10,default_factory=<dataclasses._MISSING_TYPE object>,init=True,repr=True,hash=None,compare=True,metadata=mappingproxy({'help': 'Save model every n epochs.'}),kw_only=False,_field_type=_FIELD), 'use_auth_token': Field(name='use_auth_token',type=<class 'bool'>,default=False,default_factory=<dataclasses._MISSING_TYPE object>,init=True,repr=True,hash=None,compare=True,metadata=mappingproxy({'help': 'Use the token generated when using huggingface-hub (necessary to use this script with private models).'}),kw_only=False,_field_type=_FIELD), 'use_ema': Field(name='use_ema',type=<class 'bool'>,default=True,default_factory=<dataclasses._MISSING_TYPE object>,init=True,repr=True,hash=None,compare=True,metadata=mappingproxy({'help': 'Use ema.'}),kw_only=False,_field_type=_FIELD)}

__dataclass_params__ = _DataclassParams(init=True,repr=True,eq=True,order=False,unsafe_hash=False,frozen=False)

__doc__ = 'Training arguments related to Diffusion trainer.'

__eq__(other)

Return self==value.

__hash__ = None

__init__(local_rank=-1, output_dir='ddpm-cifar10-32', logging_dir='logs/', overwrite_output_dir=False, cache_dir='.cache/', save_images_epochs=10, save_model_epochs=10, use_ema=True, ema_inv_gamma=1.0, ema_power=0.75, ema_max_decay=0.9999, mixed_precision='no', use_auth_token=False, dummy_training=False, is_sampling=True)

__match_args__ = ('local_rank', 'output_dir', 'logging_dir', 'overwrite_output_dir', 'cache_dir', 'save_images_epochs', 'save_model_epochs', 'use_ema', 'ema_inv_gamma', 'ema_power', 'ema_max_decay', 'mixed_precision', 'use_auth_token', 'dummy_training', 'is_sampling')

__module__ = 'gt4sd.training_pipelines.diffusion.core'

__repr__()

Return repr(self).

class DiffusionSavingArguments

Bases: TrainingPipelineArguments

Saving arguments related to Diffusion trainer.

__name__ = 'DiffusionSavingArguments'

__annotations__ = {}

__dataclass_fields__ = {}

__dataclass_params__ = _DataclassParams(init=True,repr=True,eq=True,order=False,unsafe_hash=False,frozen=False)

__doc__ = 'Saving arguments related to Diffusion trainer.'

__eq__(other)

Return self==value.

__hash__ = None

__init__()

__match_args__ = ()

__module__ = 'gt4sd.training_pipelines.diffusion.core'

__repr__()

Return repr(self).