| """IRIS Training utilities: synthetic dataset and scheduler.""" |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from torch.utils.data import Dataset, DataLoader |
| import math |
| import time |
| import os |
|
|
| from .model import IRIS |
| from .flow_matching import flow_matching_loss, euler_sample, DCAE_F32C32_SCALE |
|
|
|
|
| class SyntheticLatentDataset(Dataset): |
| """Generates synthetic latent/text pairs for testing training stability.""" |
| def __init__(self, num_samples=10000, latent_channels=32, latent_size=16, text_dim=512, text_length=32, seed=42): |
| self.num_samples = num_samples |
| gen = torch.Generator().manual_seed(seed) |
| self.latents = torch.randn(num_samples, latent_channels, latent_size, latent_size, generator=gen) * 2.5 |
| self.text_embeds = torch.randn(num_samples, text_length, text_dim, generator=gen) |
| self.text_embeds = F.normalize(self.text_embeds, dim=-1) * math.sqrt(text_dim) |
|
|
| def __len__(self): |
| return self.num_samples |
|
|
| def __getitem__(self, idx): |
| return {"latent": self.latents[idx], "text_embed": self.text_embeds[idx]} |
|
|
|
|
| class CosineWarmupScheduler: |
| """Cosine decay with linear warmup.""" |
| def __init__(self, optimizer, warmup_steps, total_steps, min_lr_ratio=0.1): |
| self.optimizer = optimizer |
| self.warmup_steps = warmup_steps |
| self.total_steps = total_steps |
| self.min_lr_ratio = min_lr_ratio |
| self.base_lrs = [pg["lr"] for pg in optimizer.param_groups] |
| self.step_count = 0 |
|
|
| def step(self): |
| self.step_count += 1 |
| if self.step_count <= self.warmup_steps: |
| scale = self.step_count / max(1, self.warmup_steps) |
| else: |
| progress = (self.step_count - self.warmup_steps) / max(1, self.total_steps - self.warmup_steps) |
| scale = self.min_lr_ratio + (1 - self.min_lr_ratio) * 0.5 * (1 + math.cos(math.pi * progress)) |
| for pg, base_lr in zip(self.optimizer.param_groups, self.base_lrs): |
| pg["lr"] = base_lr * scale |
|
|
| def get_lr(self): |
| return [pg["lr"] for pg in self.optimizer.param_groups] |
|
|
