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, load_from_disk
from torch.utils.data import DataLoader, TensorDataset
from torch.utils.tensorboard import SummaryWriter
from safetensors.torch import save_file, load_file
import os, time
from models import AsymmetricResidualUDiT, xATGLU
from torch.cuda.amp import autocast

from torch.optim.lr_scheduler import CosineAnnealingLR
from torch.distributions import Normal
from schedulefree import AdamWScheduleFree
from distributed_shampoo import AdamGraftingConfig, DistributedShampoo

# Changes
# MAE replace MSE
# Larger shampoo preconditioner step for stability
# Larger shampoo preconditioner dim 1024 -> 2048
# Commented out norm.

def preload_dataset(image_size=256, device="cuda", max_images=50000):
    """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")
    #dataset = load_from_disk("./new_dataset")
    
    transform = transforms.Compose([
        transforms.ToTensor(),
        #transforms.Pad((35, 0), fill=0),  # Add 35 pixels on each side horizontally (70 total to get from 186 to 256)
        transforms.Resize((256, 256), antialias=True),
        transforms.Lambda(lambda x: (x * 2) - 1)  # Scale to [-1, 1]
    ])

    all_images = []
    
    for i, example in enumerate(dataset):
        if max_images and i >= max_images:
            break
            
        img_tensor = transform(example['image'])
        
        all_images.extend([
            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]
        # Uniform Dist 0..1 -- t ~ U[0, 1]
        t = torch.rand(batch_size, 1, 1, 1, device=device)

        # Sample noise -- x0 ~ N[0, I]
        x0 = torch.randn_like(x1)
        
        # Compute OT conditional flow matching path interpolation
        
        # My understanding of this process -- We start at some random time t (Per sample)
        # We have a pure noise value at x0, which is a totally destroyed signal. 
        # We have the actual image as x1 which is a perfect signal.
        # We are going to destroy an amount of the image equal to t% of the signal. So if t is 0.3 we're destroying about 30% of the signal(image)
        # The final x_t represents our combined noisy singal, you can imagine 30% random noise overlayed onto the normal image.
        # We calculate the shortest path between x0 and x1, a straight line segment (lets call it a displacement vector) in their respective space, conditioned on the timestep.
        # We then try to predict the displacement vector where we provide our partially noisy signal and our conditioning timestep
        # We check the prediction against the real displacement vector we calculated to see how good the prediction was. Then we back propogate, baby.

        sigma_t = 1 - (1 - self.sigma_min) * t # As t increases this value decreases. This is almost 1 - t
        mu_t = t * x1 # As t increases this increases.
        x_t = sigma_t * x0 + mu_t # This is essentially a mixture of noise and signal ((1-t) * x0) + ((t) * x1)
         
        # Compute target
        target = x1 - (1 - self.sigma_min) * x0 # This is the target displacement vector (direction and magnitude) that we need to travel from x0 to x1.
        v_t = model(x_t, t) # v_t is our displacement vector prediction
        
        # Magnitude-corrected MSE
        # The 69 factor helps with very small gradients, as this loss tends to be b/w [0..1], this rescales to something more like [0..69] 
        # Other values like 420 might lead to numerical instability if the loss is too large.
        loss = F.mse_loss(v_t, target)*69 # Compare the displacement vector the network predicted to the actual displacement we calculated as mean absolute error.
        
        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)
    
    # Gradient logging
    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=1000, 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=16
    ).to(device).to(torch.float32)
    model.train()
    count_parameters(model)
    
    # optimizer = AdamWScheduleFree(
    #     model.parameters(),
    #     lr=4e-5,
    #     warmup_steps=100
    # )
    # optimizer.train()
    
    optimizer = DistributedShampoo(
        model.parameters(),
        lr=0.001,
        betas=(0.9, 0.999),
        epsilon=1e-10,
        weight_decay=1e-05,
        max_preconditioner_dim=2048,
        precondition_frequency=100,
        start_preconditioning_step=250,
        use_decoupled_weight_decay=False,
        grafting_config=AdamGraftingConfig(
            beta2=0.999,
            epsilon=1e-10,
        ),
    )
    
    scaler = torch.amp.GradScaler("cuda")
    
    scheduler = CosineAnnealingLR(
        optimizer,
        T_max=num_epochs,
        eta_min=1e-5 
    )
    
    current_batch_sizes = initial_batch_sizes.copy()
    next_drop_epoch = epoch_batch_drop_at
    interval_multiplier = 2
    
    torch.set_float32_matmul_precision('high')
    # torch.backends.cudnn.benchmark = True
    # torch.backends.cuda.matmul.allow_fp16_accumulation = True
    
    model = torch.compile(
        model,
        backend='inductor',
        dynamic=False,
        fullgraph=True,
        options={
            "epilogue_fusion": True,
            "max_autotune": True,
            "cuda.use_fast_math": True,
        }
    )
    
    flow_transport = OptimalTransportLinearFlowGenerator(sigma_min=0.001)
    
    current_batch_size = current_batch_sizes[-1]
    dataloader = DataLoader(dataset, batch_size=current_batch_size, shuffle=True)
    
    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 False:
            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})")
            
        curr_lr = optimizer.param_groups[0]['lr']
        
        for batch_idx, batch in enumerate(dataloader):
            optimizer.zero_grad()
            with torch.autocast(device_type='cuda', dtype=dtype):
                x1 = batch[0]
                batch_size = x1.shape[0]

                # x1 shape: B, C, H, W
                loss = flow_transport.loss(model, x1, device)
                
            scaler.scale(loss).backward()
            scaler.unscale_(optimizer)
            #torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            scaler.step(optimizer)
            scaler.update()
            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:.2e}")

        write_logs(writer, model, avg_loss, batch_idx, epoch, epoch_time, current_batch_size, curr_lr)
        if (epoch + 1) % 10 == 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) % 50 == 0:
            save_checkpoint(model, optimizer, f"step_{epoch}.safetensors")
            
        scheduler.step()

    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=[16,32,64],
        epoch_batch_drop_at=100,
        dtype=torch.bfloat16
    )
    
    print("Training complete! Samples saved in 'samples' directory")