Spaces:
Running on Zero
Running on Zero
| # Copyright (c) 2025 Hansheng Chen | |
| import os | |
| import argparse | |
| import tqdm | |
| import numpy as np | |
| import matplotlib.pyplot as plt | |
| import matplotlib.cm as cm | |
| import torch | |
| torch.backends.cuda.matmul.allow_tf32 = True | |
| torch.backends.cudnn.allow_tf32 = True | |
| from lakonlab.models import GMFlowMLP2DDenoiser, GMFlow | |
| from lakonlab.datasets import CheckerboardData | |
| D = 2 | |
| EPS = 1e-4 | |
| T_SCALE = 1000 | |
| N_TEST_SAMPLES = 1e6 | |
| def parse_args(): | |
| parser = argparse.ArgumentParser( | |
| description='A minimal GMFlow trainer using the 2D checkerboard dataset (without transition loss and EMA).') | |
| parser.add_argument('-k', type=int, default=32, help='number of Gasussian components') | |
| parser.add_argument('--batch-size', type=int, default=4096, help='batch size') | |
| parser.add_argument('--num-iters', type=int, default=50000, help='number of iterations') | |
| parser.add_argument('--lr', type=float, default=2e-4, help='learning rate') | |
| parser.add_argument('--nfe', type=int, default=8, help='number of sampling steps') | |
| parser.add_argument('--out', type=str, default='gmflow_toymodel.png', help='output file path') | |
| args = parser.parse_args() | |
| return args | |
| def gm_kl_loss(gm, sample, eps=1e-4): | |
| """ | |
| Gaussian mixture KL divergence loss (without constant terms), a.k.a. GM NLL loss. | |
| Args: | |
| gm (dict): | |
| means (torch.Tensor): (bs, num_gaussians, D) | |
| logstds (torch.Tensor): (bs, 1, 1) | |
| logweights (torch.Tensor): (bs, num_gaussians, 1) | |
| sample (torch.Tensor): (bs, D) | |
| Returns: | |
| torch.Tensor: (bs, ) | |
| """ | |
| means = gm['means'] | |
| logstds = gm['logstds'] | |
| logweights = gm['logweights'] | |
| inverse_stds = torch.exp(-logstds).clamp(max=1 / eps) | |
| diff_weighted = (sample.unsqueeze(-2) - means) * inverse_stds # (bs, num_gaussians, D) | |
| gaussian_ll = (-0.5 * diff_weighted.square() - logstds).sum(dim=-1) # (bs, num_gaussians) | |
| gm_nll = -torch.logsumexp(gaussian_ll + logweights.squeeze(-1), dim=-1) # (bs, ) | |
| return gm_nll | |
| def main(): | |
| args = parse_args() | |
| num_gaussians = args.k | |
| batch_size = args.batch_size | |
| num_iters = args.num_iters | |
| lr = args.lr | |
| num_steps = args.nfe | |
| out_path = args.out | |
| device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') | |
| denoiser = GMFlowMLP2DDenoiser(num_gaussians=num_gaussians).to(device) | |
| dataset = CheckerboardData(scale=4) | |
| sample_batches = dataset.samples.to(device).split(batch_size, dim=0) | |
| optimizer = torch.optim.Adam(denoiser.parameters(), lr=lr) | |
| for i in range(num_iters): | |
| x_0 = sample_batches[i % len(sample_batches)] | |
| optimizer.zero_grad() | |
| t = torch.rand(x_0.size(0), device=device).clamp(min=EPS) | |
| noise = torch.randn_like(x_0) | |
| sigma = t | |
| alpha = 1 - sigma | |
| x_t = alpha.unsqueeze(-1) * x_0 + sigma.unsqueeze(-1) * noise | |
| u = noise - x_0 # equal to (x_t - x_0) / sigma | |
| u_gm = denoiser(x_t, t * T_SCALE) | |
| loss = gm_kl_loss(u_gm, u) | |
| loss.mean().backward() | |
| optimizer.step() | |
| if i % 1000 == 0: | |
| print(f'Iter {i}, loss: {loss.mean().item()}') | |
| print('Training finished. Starting inference...') | |
| torch.set_grad_enabled(False) | |
| model = GMFlow( | |
| denoising=denoiser, | |
| num_timesteps=T_SCALE, | |
| test_cfg=dict( # use 2nd-order GM-SDE solver | |
| output_mode='sample', | |
| sampler='FlowSDE', | |
| num_timesteps=num_steps, | |
| order=2) | |
| ).eval() | |
| samples = [] | |
| for _ in tqdm.tqdm(range(int(N_TEST_SAMPLES // batch_size))): | |
| noise = torch.randn((batch_size, D, 1, 1), device=device) | |
| samples.append(model.forward_test(noise=noise).reshape(batch_size, D).cpu().numpy()) | |
| samples = np.concatenate(samples, axis=0) | |
| histo, _, _ = np.histogram2d( | |
| samples[:, 0], samples[:, 1], bins=200, range=[[-4.2, 4.2], [-4.2, 4.2]]) | |
| histo_image = (histo.T[::-1] / 160).clip(0, 1) | |
| histo_image = cm.viridis(histo_image) | |
| histo_image = np.round(histo_image * 255).clip(min=0, max=255).astype(np.uint8) | |
| out_path = os.path.abspath(out_path) | |
| out_dir = os.path.dirname(out_path) | |
| os.makedirs(out_dir, exist_ok=True) | |
| plt.imsave(out_path, histo_image) | |
| print(f'Sample histogram saved to {out_path}.') | |
| if __name__ == '__main__': | |
| main() | |