def init(self):
# Create $\color{cyan}{\epsilon_\theta}(x_t, t)$ model
self.eps_model = UNet(
image_channels=self.image_channels,
n_channels=self.n_channels,
ch_mults=self.channel_multipliers,
is_attn=self.is_attention,
).to(self.device)
# Create [DDPM class](index.html)
self.diffusion = DenoiseDiffusion(
eps_model=self.eps_model,
n_steps=self.n_steps,
device=self.device,
)
# Create dataloader
self.data_loader = torch.utils.data.DataLoader(self.dataset,
self.batch_size,
shuffle=True,
pin_memory=True)
# Create optimizer
self.optimizer = torch.optim.Adam(self.eps_model.parameters(),
lr=self.learning_rate)
# Image logging
tracker.set_image("sample", True)
def init(self):
"""
### Initialize
"""
# Create dataset
dataset = Dataset(self.dataset_path, self.image_size)
# Create data loader
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=self.batch_size,
num_workers=8,
shuffle=True,
drop_last=True,
pin_memory=True)
# Continuous [cyclic loader](../../utils.html#cycle_dataloader)
self.loader = cycle_dataloader(dataloader)
# $\log_2$ of image resolution
log_resolution = int(math.log2(self.image_size))
# Create discriminator and generator
self.discriminator = Discriminator(log_resolution).to(self.device)
self.generator = Generator(log_resolution,
self.d_latent).to(self.device)
# Get number of generator blocks for creating style and noise inputs
self.n_gen_blocks = self.generator.n_blocks
# Create mapping network
self.mapping_network = MappingNetwork(
self.d_latent, self.mapping_network_layers).to(self.device)
# Create path length penalty loss
self.path_length_penalty = PathLengthPenalty(0.99).to(self.device)
# Add model hooks to monitor layer outputs
if self.log_layer_outputs:
hook_model_outputs(self.mode, self.discriminator, 'discriminator')
hook_model_outputs(self.mode, self.generator, 'generator')
hook_model_outputs(self.mode, self.mapping_network,
'mapping_network')
# Discriminator and generator losses
self.discriminator_loss = DiscriminatorLoss().to(self.device)
self.generator_loss = GeneratorLoss().to(self.device)
# Create optimizers
self.discriminator_optimizer = torch.optim.Adam(
self.discriminator.parameters(),
lr=self.learning_rate,
betas=self.adam_betas)
self.generator_optimizer = torch.optim.Adam(
self.generator.parameters(),
lr=self.learning_rate,
betas=self.adam_betas)
self.mapping_network_optimizer = torch.optim.Adam(
self.mapping_network.parameters(),
lr=self.mapping_network_learning_rate,
betas=self.adam_betas)
# Set tracker configurations
tracker.set_image("generated", True)
def init(self):
"""
Initializations
"""
self.state_modules = []
hook_model_outputs(self.mode, self.generator, 'generator')
hook_model_outputs(self.mode, self.discriminator, 'discriminator')
tracker.set_scalar("loss.generator.*", True)
tracker.set_scalar("loss.discriminator.*", True)
tracker.set_image("generated", True, 1 / 100)
def init(self):
self.state_modules = []
self.generator = Generator().to(self.device)
self.discriminator = Discriminator().to(self.device)
self.generator_loss = GeneratorLogitsLoss(self.label_smoothing).to(self.device)
self.discriminator_loss = DiscriminatorLogitsLoss(self.label_smoothing).to(self.device)
hook_model_outputs(self.mode, self.generator, 'generator')
hook_model_outputs(self.mode, self.discriminator, 'discriminator')
tracker.set_scalar("loss.generator.*", True)
tracker.set_scalar("loss.discriminator.*", True)
tracker.set_image("generated", True, 1 / 100)
def __init__(self, *, discriminator: Module, generator: Module,
discriminator_optimizer: Optional[torch.optim.Adam],
generator_optimizer: Optional[torch.optim.Adam],
discriminator_loss: DiscriminatorLogitsLoss,
generator_loss: GeneratorLogitsLoss, discriminator_k: int):
self.discriminator_k = discriminator_k
self.generator = generator
self.discriminator = discriminator
self.generator_loss = generator_loss
self.discriminator_loss = discriminator_loss
self.generator_optimizer = generator_optimizer
self.discriminator_optimizer = discriminator_optimizer
hook_model_outputs(self.generator, 'generator')
hook_model_outputs(self.discriminator, 'discriminator')
tracker.set_scalar("loss.generator.*", True)
tracker.set_scalar("loss.discriminator.*", True)
tracker.set_image("generated", True, 1 / 100)
def __call__(self):
tracker.set_image('generated')
for _ in self.loop:
self._train()