Upload iris/train.py

iris/train.py

@@ -0,0 +1,52 @@
+"""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]