smalldiffusion 实战示例

import numpy as np
from torch.utils.data import DataLoader
from smalldiffusion import Swissroll, TimeInputMLP, ScheduleLogLinear, training_loop, samples

# 1. 数据
dataset  = Swissroll(np.pi/2, 5*np.pi, 100)
loader   = DataLoader(dataset, batch_size=2048)

# 2. 模型
model    = TimeInputMLP(hidden_dims=(16,128,128,128,128,16))

# 3. 调度
schedule = ScheduleLogLinear(N=200, sigma_min=0.005, sigma_max=10)

# 4. 训练
trainer  = training_loop(loader, model, schedule, epochs=15000)
losses   = [ns.loss.item() for ns in trainer]

# 5. 采样
*xt, x0  = samples(model, schedule.sample_sigmas(20), gam=2)
# x0.shape: (1, 2) — 一个生成的 2D 点
# xt: 20 个中间步的结果列表

import torch
import numpy as np
from torch.utils.data import DataLoader
from smalldiffusion import (
    TreeDataset, ConditionalMLP, ScheduleLogLinear,
    training_loop, samples
)

# 数据：4 分支、3 层深度的树
dataset = TreeDataset(branching_factor=4, depth=3)
loader = DataLoader(dataset, batch_size=2048, shuffle=True)

# 条件模型
model = ConditionalMLP(
    dim=2,
    hidden_dims=(16, 128, 256, 128, 16),
    cond_dim=4,
    num_classes=dataset.total_leaves,  # 64 类
    dropout_prob=0.1
)

# 训练（conditional=True）
schedule = ScheduleLogLinear(N=200, sigma_min=0.01, sigma_max=10)
trainer = training_loop(loader, model, schedule, epochs=15000, conditional=True)
losses = [ns.loss.item() for ns in trainer]

# 条件采样 + CFG
N_sample = 64
cond = torch.arange(dataset.total_leaves)  # 每个类别各生成一个样本
*xt, x0 = samples(model, schedule.sample_sigmas(20),
                   gam=2, batchsize=N_sample,
                   cond=cond, cfg_scale=4.0)

from smalldiffusion import (
    ScheduleDDPM, samples, training_loop, MappedDataset, DiT,
    img_train_transform, img_normalize
)
from torch_ema import ExponentialMovingAverage as EMA

# 数据：丢弃标签
dataset = MappedDataset(
    FashionMNIST('datasets', train=True, download=True, transform=img_train_transform),
    lambda x: x[0]
)
loader = DataLoader(dataset, batch_size=1024, shuffle=True)

# 模型
schedule = ScheduleDDPM(beta_start=0.0001, beta_end=0.02, N=1000)
model = DiT(in_dim=28, channels=1, patch_size=2, depth=6,
            head_dim=32, num_heads=6, mlp_ratio=4.0)

# 训练 + EMA
ema = EMA(model.parameters(), decay=0.99)
ema.to(accelerator.device)
for ns in training_loop(loader, model, schedule, epochs=300, lr=1e-3, accelerator=a):
    ns.pbar.set_description(f'Loss={ns.loss.item():.5}')
    ema.update()

# 采样（使用 EMA 参数）
with ema.average_parameters():
    *xt, x0 = samples(model, schedule.sample_sigmas(20), gam=1.6, batchsize=64)
    save_image(img_normalize(make_grid(x0)), 'samples.png')

# 不丢弃标签
dataset = FashionMNIST('datasets', train=True, download=True, transform=img_train_transform)

# 添加条件嵌入
model = DiT(in_dim=28, channels=1, patch_size=2, depth=6,
            head_dim=32, num_heads=6, mlp_ratio=4.0,
            cond_embed=CondEmbedderLabel(32*6, 10, 0.1))  # 10 类，10% dropout

# 条件训练
for ns in training_loop(loader, model, schedule, epochs=300, conditional=True, ...):
    ...

# 条件采样
cond = torch.tensor([i % 10 for i in range(40)])  # 每类 4 个样本
*xt, x0 = samples(model, schedule.sample_sigmas(20), gam=1.6, batchsize=40, cond=cond)

from smalldiffusion import Unet, Scaled, ScheduleLogLinear, ScheduleSigmoid

# 训练和采样使用不同调度
train_schedule = ScheduleSigmoid(N=1000)
model = Scaled(Unet)(32, 3, 3, ch=128, ch_mult=(1, 2, 2, 2), attn_resolutions=(16,))

# 训练...

# 采样使用不同调度
sample_schedule = ScheduleLogLinear(sigma_min=0.01, sigma_max=35, N=1000)
*xt, x0 = samples(model, sample_schedule.sample_sigmas(10), gam=2.1, batchsize=64)