Upload train.py

train.py

@@ -0,0 +1,335 @@
+"""
+LiquidFlow Trainer — Complete training pipeline.
+
+Usage:
+    python train.py --dataset cifar10 --image_size 128 --variant small --batch_size 32 --epochs 100
+
+Features:
+- Automatic VAE loading (TAESD by default)
+- Physics-informed regularization
+- Mixed precision training (AMP)
+- Checkpoint saving
+- Sample generation during training
+- Colab/Kaggle compatible (T4 GPU, 15GB VRAM)
+
+Requirements:
+    pip install torch torchvision diffusers tqdm pillow numpy
+"""
+
+import os
+import sys
+import math
+import argparse
+import json
+from datetime import datetime
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+from torchvision import datasets, transforms
+from torchvision.utils import save_image
+import numpy as np
+from tqdm import tqdm
+
+# Add parent to path
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from liquid_flow.generator import LiquidFlowGenerator, create_liquidflow
+from liquid_flow.vae_wrapper import TAESDWrapper
+
+
+def get_dataloader(dataset_name, image_size, batch_size, data_dir='./data'):
+    """Get training dataloader for common datasets."""
+    transform = transforms.Compose([
+        transforms.Resize((image_size, image_size)),
+        transforms.ToTensor(),
+        transforms.Normalize([0.5], [0.5]),  # [-1, 1]
+    ])
+    
+    if dataset_name == 'cifar10':
+        dataset = datasets.CIFAR10(
+            root=data_dir, train=True, download=True, transform=transform
+        )
+    elif dataset_name == 'cifar100':
+        dataset = datasets.CIFAR100(
+            root=data_dir, train=True, download=True, transform=transform
+        )
+    elif dataset_name == 'stl10':
+        dataset = datasets.STL10(
+            root=data_dir, split='train', download=True, transform=transform
+        )
+    elif dataset_name == 'celeba':
+        dataset = datasets.CelebA(
+            root=data_dir, split='train', download=True, transform=transform
+        )
+    elif dataset_name == 'lsun':
+        dataset = datasets.LSUN(
+            root=data_dir, classes='bedroom_train', transform=transform
+        )
+    elif dataset_name == 'imagenet':
+        transform = transforms.Compose([
+            transforms.Resize((image_size, image_size)),
+            transforms.RandomCrop(image_size),
+            transforms.RandomHorizontalFlip(),
+            transforms.ToTensor(),
+            transforms.Normalize([0.5], [0.5]),
+        ])
+        dataset = datasets.ImageFolder(
+            root=f'{data_dir}/imagenet/train', transform=transform
+        )
+    else:
+        raise ValueError(f"Unknown dataset: {dataset_name}")
+    
+    dataloader = DataLoader(
+        dataset,
+        batch_size=batch_size,
+        shuffle=True,
+        num_workers=min(4, os.cpu_count() or 1),
+        pin_memory=True,
+        drop_last=True,
+    )
+    
+    return dataloader
+
+
+def train(args):
+    """Main training loop."""
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    print(f"Using device: {device}")
+    
+    # Create output directory
+    os.makedirs(args.output_dir, exist_ok=True)
+    os.makedirs(os.path.join(args.output_dir, 'samples'), exist_ok=True)
+    os.makedirs(os.path.join(args.output_dir, 'checkpoints'), exist_ok=True)
+    
+    # Load VAE
+    print("Loading VAE...")
+    vae = TAESDWrapper.load(device)
+    print(f"VAE loaded. Latent size: {args.image_size // 8}x{args.image_size // 8}")
+    
+    # Create model
+    print(f"Creating LiquidFlow model (variant={args.variant})...")
+    model = create_liquidflow(
+        variant=args.variant,
+        image_size=args.image_size,
+    )
+    model = model.to(device)
+    
+    n_params = model.count_parameters()
+    print(f"Model parameters: {n_params:,} (~{n_params/1e6:.1f}M)")
+    
+    # Calculate memory estimate
+    latent_h = latent_w = args.image_size // 8
+    mem_per_sample = latent_h * latent_w * 4 * 4 / (1024**2)  # in MB
+    print(f"Estimated memory per sample: {mem_per_sample:.1f} MB")
+    print(f"Estimated batch memory: {mem_per_sample * args.batch_size:.1f} MB")
+    
+    # Dataset
+    print(f"Loading dataset: {args.dataset}")
+    dataloader = get_dataloader(args.dataset, args.image_size, args.batch_size, args.data_dir)
+    print(f"Dataset size: {len(dataloader.dataset)} images, {len(dataloader)} batches")
+    
+    # Optimizer (AdamW, following DiT/DiMSUM convention)
+    optimizer = torch.optim.AdamW(
+        model.parameters(),
+        lr=args.lr,
+        betas=(0.9, 0.999),
+        weight_decay=args.weight_decay,
+    )
+    
+    # Learning rate scheduler
+    if args.lr_schedule == 'cosine':
+        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
+            optimizer, T_max=args.epochs * len(dataloader)
+        )
+    elif args.lr_schedule == 'cosine_restart':
+        scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(
+            optimizer, T_0=args.epochs * len(dataloader) // 3,
+        )
+    else:
+        scheduler = None
+    
+    # AMP
+    use_amp = args.amp and device.type == 'cuda'
+    scaler = torch.cuda.amp.GradScaler() if use_amp else None
+    
+    # Fixed noise for sample generation (track progress)
+    sample_noise = torch.randn(16, 4, args.image_size // 8, args.image_size // 8, device=device)
+    
+    # Training state
+    global_step = 0
+    best_loss = float('inf')
+    
+    print(f"\n{'='*60}")
+    print(f"Starting training: {args.epochs} epochs, {args.batch_size} batch size")
+    print(f"LR: {args.lr}, Weight Decay: {args.weight_decay}")
+    print(f"AMP: {use_amp}, LR Schedule: {args.lr_schedule}")
+    print(f"{'='*60}\n")
+    
+    for epoch in range(args.epochs):
+        model.train()
+        epoch_losses = {'total': 0, 'diffusion': 0, 'physics': 0}
+        
+        pbar = tqdm(dataloader, desc=f"Epoch {epoch+1}/{args.epochs}")
+        
+        for batch_idx, (images, _) in enumerate(pbar):
+            images = images.to(device)
+            
+            # Encode to latent space
+            with torch.no_grad():
+                latents = TAESDWrapper.encode(vae, images)
+            
+            # Training step
+            loss_dict = model.training_step(latents, optimizer, scaler, use_amp)
+            
+            # Update scheduler
+            if scheduler is not None:
+                scheduler.step()
+            
+            # Track losses
+            for k in epoch_losses:
+                epoch_losses[k] += loss_dict.get(k, 0)
+            
+            global_step += 1
+            
+            # Update progress bar
+            pbar.set_postfix({
+                'loss': f"{loss_dict.get('total', 0):.4f}",
+                'diff': f"{loss_dict.get('diffusion', 0):.4f}",
+                'phys': f"{loss_dict.get('physics', 0):.4f}",
+                'lr': f"{optimizer.param_groups[0]['lr']:.2e}",
+            })
+        
+        # Epoch summary
+        n_batches = len(dataloader)
+        avg_losses = {k: v / n_batches for k, v in epoch_losses.items()}
+        
+        print(f"\nEpoch {epoch+1} Summary:")
+        print(f"  Total Loss:     {avg_losses['total']:.4f}")
+        print(f"  Diffusion Loss: {avg_losses['diffusion']:.4f}")
+        print(f"  Physics Loss:   {avg_losses['physics']:.4f}")
+        
+        # Generate samples
+        if (epoch + 1) % args.sample_every == 0 or epoch == args.epochs - 1:
+            print(f"Generating samples...")
+            model.eval()
+            
+            with torch.no_grad():
+                # DDIM sampling
+                latents_gen = model.sample(
+                    batch_size=16,
+                    steps=args.sample_steps,
+                    ddim=True,
+                    progress=False,
+                )
+                images_gen = TAESDWrapper.decode(vae, latents_gen)
+                
+                # Also generate from fixed noise for tracking
+                t_fixed = torch.full((16,), 0, device=device, dtype=torch.long)
+                # Quick DDIM from fixed noise
+                x_fixed = sample_noise.clone()
+                skip = 1000 // args.sample_steps
+                for i in reversed(range(0, 1000, skip)):
+                    t = torch.full((16,), i, device=device, dtype=torch.long)
+                    noise_pred = model(x_fixed, t)
+                    alpha_bar = model.alphas_cumprod[i]
+                    alpha_bar_prev = model.alphas_cumprod[i - skip] if i >= skip else torch.tensor(1.0, device=device)
+                    x0_pred = (x_fixed - torch.sqrt(1 - alpha_bar) * noise_pred) / torch.sqrt(alpha_bar)
+                    x0_pred = torch.clamp(x0_pred, -1, 1)
+                    x_fixed = torch.sqrt(alpha_bar_prev) * x0_pred + torch.sqrt(1 - alpha_bar_prev) * torch.randn_like(x_fixed)
+                
+                images_fixed = TAESDWrapper.decode(vae, x_fixed)
+            
+            # Save samples
+            sample_path = os.path.join(args.output_dir, 'samples', f'epoch_{epoch+1:03d}.png')
+            save_image(images_gen, sample_path, nrow=4, normalize=True, value_range=(-1, 1))
+            
+            fixed_path = os.path.join(args.output_dir, 'samples', f'fixed_{epoch+1:03d}.png')
+            save_image(images_fixed, fixed_path, nrow=4, normalize=True, value_range=(-1, 1))
+            
+            print(f"  Samples saved to {sample_path}")
+        
+        # Save checkpoint
+        if (epoch + 1) % args.save_every == 0 or epoch == args.epochs - 1:
+            checkpoint_path = os.path.join(args.output_dir, 'checkpoints', f'epoch_{epoch+1:03d}.pt')
+            torch.save({
+                'epoch': epoch + 1,
+                'global_step': global_step,
+                'model_state_dict': model.state_dict(),
+                'optimizer_state_dict': optimizer.state_dict(),
+                'loss': avg_losses['total'],
+                'args': vars(args),
+            }, checkpoint_path)
+            print(f"  Checkpoint saved to {checkpoint_path}")
+            
+            # Save best model
+            if avg_losses['total'] < best_loss:
+                best_loss = avg_losses['total']
+                best_path = os.path.join(args.output_dir, 'checkpoints', 'best_model.pt')
+                torch.save(model.state_dict(), best_path)
+                print(f"  Best model saved (loss={best_loss:.4f})")
+        
+        print()
+    
+    print(f"\n{'='*60}")
+    print(f"Training complete!")
+    print(f"Best loss: {best_loss:.4f}")
+    print(f"Model saved to: {args.output_dir}/checkpoints/")
+    print(f"{'='*60}")
+    
+    return model
+
+
+def main():
+    parser = argparse.ArgumentParser(description='LiquidFlow Generator Training')
+    
+    # Dataset
+    parser.add_argument('--dataset', type=str, default='cifar10',
+                        choices=['cifar10', 'cifar100', 'stl10', 'celeba', 'lsun', 'imagenet'],
+                        help='Training dataset')
+    parser.add_argument('--data_dir', type=str, default='./data',
+                        help='Data directory')
+    parser.add_argument('--image_size', type=int, default=128,
+                        choices=[64, 128, 256, 512],
+                        help='Image size (will be VAE-encoded)')
+    
+    # Model
+    parser.add_argument('--variant', type=str, default='small',
+                        choices=['tiny', 'small', 'base'],
+                        help='Model size variant')
+    
+    # Training
+    parser.add_argument('--batch_size', type=int, default=32,
+                        help='Batch size')
+    parser.add_argument('--epochs', type=int, default=100,
+                        help='Number of epochs')
+    parser.add_argument('--lr', type=float, default=2e-4,
+                        help='Learning rate')
+    parser.add_argument('--weight_decay', type=float, default=1e-4,
+                        help='Weight decay')
+    parser.add_argument('--lr_schedule', type=str, default='cosine',
+                        choices=['cosine', 'cosine_restart', 'none'],
+                        help='LR schedule')
+    parser.add_argument('--amp', action='store_true', default=True,
+                        help='Use automatic mixed precision')
+    
+    # Generation
+    parser.add_argument('--sample_every', type=int, default=5,
+                        help='Generate samples every N epochs')
+    parser.add_argument('--sample_steps', type=int, default=50,
+                        help='DDIM sampling steps')
+    
+    # IO
+    parser.add_argument('--output_dir', type=str, default='./outputs',
+                        help='Output directory')
+    parser.add_argument('--save_every', type=int, default=10,
+                        help='Save checkpoint every N epochs')
+    
+    args = parser.parse_args()
+    train(args)
+
+
+if __name__ == '__main__':
+    main()