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")