train.py

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torchvision.transforms as transforms
import torchvision.utils as vutils
from datasets import load_dataset
from torch.utils.data import DataLoader, TensorDataset
from schedulefree import AdamWScheduleFree
from torch.utils.tensorboard import SummaryWriter
from safetensors.torch import save_file, load_file
import os, time
from models import AsymmetricResidualUDiT
from torch.cuda.amp import autocast

def preload_dataset(image_size=256, device="cuda"):
    """Preload and cache the entire dataset in GPU memory"""
    print("Loading and preprocessing dataset...")
    #dataset = load_dataset("jiovine/pixel-art-nouns-2k", split="train")
    dataset = load_dataset("reach-vb/pokemon-blip-captions", split="train")
    
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Resize((image_size, image_size), antialias=True),
        transforms.Lambda(lambda x: (x * 2) - 1)  # Scale to [-1, 1]
    ])
    
    all_images = []
    for example in dataset:
        img_tensor = transform(example['image'])
        all_images.append(img_tensor)
        
    # Stack entire dataset onto gpu
    images_tensor = torch.stack(all_images).to(device)
    print(f"Dataset loaded: {images_tensor.shape} ({images_tensor.element_size() * images_tensor.nelement() / 1024/1024:.2f} MB)")
    
    return TensorDataset(images_tensor)

def count_parameters(model):
    total_params = sum(p.numel() for p in model.parameters())
    print(f'Total parameters: {total_params:,} ({total_params/1e6:.2f}M)')
    
def save_checkpoint(model, optimizer, filename="checkpoint.safetensors"):
    model_state = model.state_dict()
    save_file(model_state, filename)

def load_checkpoint(model, optimizer, filename="checkpoint.safetensors"):
    model_state = load_file(filename)
    model.load_state_dict(model_state)

# https://arxiv.org/abs/2210.02747
class OptimalTransportLinearFlowGenerator():
    def __init__(self, sigma_min=0.001):
        self.sigma_min = sigma_min
        
    def loss(self, model, x1, device):
        batch_size = x1.shape[0]
        
        # Sample t uniform in [0,1]
        t = torch.rand(batch_size, 1, 1, 1, device=device)
        
        # Sample noise
        x0 = torch.randn_like(x1)
        x1 = x1
        
        # Compute OT path interpolation (equation 22)
        sigma_t = 1 - (1 - self.sigma_min) * t
        mu_t = t * x1
        x_t = sigma_t * x0 + mu_t
        
        # Compute target (equation 23)
        target = x1 - (1 - self.sigma_min) * x0
        
        v_t = model(x_t, t)
        loss = F.mse_loss(v_t, target)
        
        return loss

def write_logs(writer, model, loss, batch_idx, epoch, epoch_time, batch_size, lr, log_gradients=True):
    """
    TensorBoard logging
    
    Args:
        writer: torch.utils.tensorboard.SummaryWriter instance
        model: torch.nn.Module - the model being trained
        loss: float or torch.Tensor - the loss value to log
        batch_idx: int - current batch index
        epoch: int - current epoch
        epoch_time: float - time taken for epoch
        batch_size: int - current batch size
        lr: float - current learning rate
        samples: Optional[torch.Tensor] - generated samples to log (only passed every 50 epochs)
        log_gradients: bool - whether to log gradient norms
    """
    total_steps = epoch * batch_idx
    
    writer.add_scalar('Loss/batch', loss, total_steps)
    writer.add_scalar('Time/epoch', epoch_time, epoch)
    writer.add_scalar('Training/batch_size', batch_size, epoch)
    writer.add_scalar('Training/learning_rate', lr, epoch)
    
    if log_gradients:
        total_norm = 0.0
        for p in model.parameters():
            if p.grad is not None:
                param_norm = p.grad.detach().data.norm(2)
                total_norm += param_norm.item() ** 2
        total_norm = total_norm ** 0.5
        writer.add_scalar('Gradients/total_norm', total_norm, total_steps)
    
def train_udit_flow(num_epochs=5000, initial_batch_sizes=[8, 16, 32, 64, 128], epoch_batch_drop_at=40, device="cuda", dtype=torch.float32):
    dataset = preload_dataset(device=device)
    temp_loader = DataLoader(dataset, batch_size=initial_batch_sizes[0], shuffle=True)
    first_batch = next(iter(temp_loader))
    image_shape = first_batch[0].shape[1:]
    
    writer = SummaryWriter('logs/current_run')
    
    model = AsymmetricResidualUDiT(
        in_channels=3,
        base_channels=128,
        num_levels=3,
        patch_size=4,
        encoder_blocks=3,
        decoder_blocks=7,
        encoder_transformer_thresh=2,
        decoder_transformer_thresh=4,
        mid_blocks=8
    ).to(device).to(dtype)
    model.train()
    
    count_parameters(model)
    optimizer = AdamWScheduleFree(
        model.parameters(),
        lr=1e-4,
        warmup_steps=100
    )
    optimizer.train()
    
    current_batch_sizes = initial_batch_sizes.copy()
    next_drop_epoch = epoch_batch_drop_at
    interval_multiplier = 2
    
    torch.set_float32_matmul_precision('high')
    model = torch.compile(
        model,
        backend='inductor',
        mode='max-autotune',
        fullgraph=True,
    )
    
    flow_transport = OptimalTransportLinearFlowGenerator(sigma_min=0.001)

    for epoch in range(num_epochs):
        epoch_start_time = time.time()
        total_loss = 0
        
        # Batch size decay logic
        # Geomtric growth, every X*N+(X-1*N+...) use the number batch size in the list.
        if epoch > 0 and epoch == next_drop_epoch and len(current_batch_sizes) > 1:
            current_batch_sizes.pop()
            next_interval = epoch_batch_drop_at * interval_multiplier
            next_drop_epoch += next_interval
            interval_multiplier += 1
            print(f"\nEpoch {epoch}: Reducing batch size to {current_batch_sizes[-1]}")
            print(f"Next drop will occur at epoch {next_drop_epoch} (interval: {next_interval})")
            
        current_batch_size = current_batch_sizes[-1]
        dataloader = DataLoader(dataset, batch_size=current_batch_size, shuffle=True)
        curr_lr = optimizer.param_groups[0]['lr']
        
        with torch.amp.autocast('cuda', dtype=dtype):
            for batch_idx, batch in enumerate(dataloader):
                x1 = batch[0]
                batch_size = x1.shape[0]
                
                loss = flow_transport.loss(model, x1, device)
                
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()
                total_loss += loss.item()
            
        avg_loss = total_loss / len(dataloader)
        
        epoch_time = time.time() - epoch_start_time
        print(f"Epoch {epoch}, Took: {epoch_time:.2f}s, Batch Size: {current_batch_size}, "
              f"Average Loss: {avg_loss:.4f}, Learning Rate: {curr_lr:.6f}")

        write_logs(writer, model, avg_loss, batch_idx, epoch, epoch_time, current_batch_size, curr_lr)
        if (epoch + 1) % 50 == 0:
            with torch.amp.autocast('cuda', dtype=dtype):
                sampling_start_time = time.time()
                samples = sample(model, device=device, dtype=dtype)
                os.makedirs("samples", exist_ok=True)
                vutils.save_image(samples, f"samples/epoch_{epoch}.png", nrow=4, padding=2)
                
                sample_time = time.time() - sampling_start_time
                print(f"Sampling took: {sample_time:.2f}s")
                
        if (epoch + 1) % 200 == 0:
            save_checkpoint(model, optimizer, f"step_{epoch}.safetensors")

    return model

def sample(model, n_samples=16, n_steps=50, image_size=256, device="cuda", sigma_min=0.001, dtype=torch.float32):
    with torch.amp.autocast('cuda', dtype=dtype):
        
        x = torch.randn(n_samples, 3, image_size, image_size, device=device)
        ts = torch.linspace(0, 1, n_steps, device=device)
        dt = 1/n_steps
        
        # Forward Euler Integration step 0..1
        with torch.no_grad():
            for i in range(len(ts)):
                t = ts[i]
                t_input = t.repeat(n_samples, 1, 1, 1)
                
                v_t = model(x, t_input)
                
                x = x + v_t * dt
    
    return x.float()

if __name__ == "__main__":
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"Using device: {device}")
    
    model = train_udit_flow(
        device=device,
        initial_batch_sizes=[8, 16],
        epoch_batch_drop_at=600,
        dtype=torch.float32
    )
    
    print("Training complete! Samples saved in 'samples' directory")