Tiny Universe, часть 6: TinyDiffusion — DDPM с нуля на CIFAR-10

О переводе

Полный перевод файла content/TinyDiffusion/readme.md проекта tiny-universe (CC BY-NC-SA 4.0 → эта страница тоже под CC BY-NC-SA 4.0). Переведены текст, формулы с пояснениями и комментарии в коде; рисунки не копируем — описываем и даём ссылки. Наши пояснения — «Примечание переводчика»; в них, в частности, разобраны две неточности кода оценки. Автор раздела в самом файле не назван.

Tiny DDPM

Как пользоваться (Usage)

pip install -r requirements.txt

• Скачайте датасет CIFAR-10 в папку datasets/, чтобы получилась структура datasets/cifar-10-batches-py.
• Обучить модель: python ddpm/train.py
• Сгенерировать изображения: python ddpm/sample.py
• Посчитать IS и FID: python ddpm/metrics.py

Напоминание: формулы DDPM

Рисунок в оригинале (fig/fig1.png) — марковская цепь DDPM: $x_T \to \dots \to x_t \to x_{t-1} \to \dots \to x_0$ , переходы назад задаёт модель $p_\theta(x_{t-1}\mid x_t)$, вперёд — зашумление $q(x_t \mid x_{t-1})$.

DDPM — генеративная модель на основе марковской цепи. Как показано на рисунке, прямой процесс (Forward Process) постепенно добавляет к данным гауссов шум, пока не останется чистый шум, а обратный процесс (Reverse Process) постепенно восстанавливает данные из шума.

Прямое зашумление устроено просто. А из формул прямого процесса нужно вывести формулу обратного, очищающего процесса — она и станет целью обучения модели. При самом обучении достаточно формулы обратного процесса.

Обозначения

$p_\theta$ — распределение вероятностей процесса очистки, которое оценивает сеть с параметрами $\theta$

$q$ — истинное распределение изображений

$x_0$ — исходное изображение

$x_t$ — изображение после $t$ шагов зашумления

Прямой процесс: зашумление

В прямом процессе к $x_{t-1}$ добавляется небольшой гауссов шум, и получается значение $x_t$ в следующий момент:

$x_t = \sqrt{1-\beta_t}\,x_{t-1} + \sqrt{\beta_t}\,\epsilon, \quad \epsilon \sim \mathcal{N}(0,I)$

Здесь $\beta_t$ — заранее заданный коэффициент дисперсии, он управляет зашумлением и выбирается вручную.

Это формула (2) статьи (там она записана как распределение):

$q(x_t \mid x_{t-1}) := \mathcal{N}\big(x_t;\ \sqrt{1-\beta_t}\,x_{t-1},\ \beta_t I\big)$

Обозначим $\alpha_t = 1-\beta_t$ и продолжим подстановку дальше; по свойству сложения гауссовых величин получаем формулу перехода сразу от $x_0$ к $x_t$:

$x_t = \sqrt{\bar{\alpha}_t}\,x_0 + \sqrt{1-\bar{\alpha}_t}\,\epsilon, \quad \epsilon \sim \mathcal{N}(0,I)$

где $\bar{\alpha}_t = \prod_{i=1}^t \alpha_i$.

Это формула (4) статьи:

$q(x_t \mid x_0) := \mathcal{N}\big(x_t;\ \sqrt{\bar{\alpha}_t}\,x_0,\ (1-\bar{\alpha}_t)I\big)$

Когда $t$ достаточно велико ($t \rightarrow T$), $x_T$ почти распределено как стандартное нормальное: $x_T \rightarrow \epsilon$. То есть после достаточного числа шагов зашумления изображение полностью превращается в случайный шум. Поэтому при генерации в качестве стартовой точки можно просто взять выборку из стандартного нормального распределения.

Примечание переводчика: формула в числах

В коде ниже $\beta_t$ растёт линейно от 0,0001 до 0,02 за $T = 1000$ шагов. Посчитаем коэффициенты формулы $x_t = \sqrt{\bar\alpha_t}\,x_0 + \sqrt{1-\bar\alpha_t}\,\epsilon$:

шаг $t$	$\bar\alpha_t$	доля картинки $\sqrt{\bar\alpha_t}$	доля шума $\sqrt{1-\bar\alpha_t}$
1	0,9999	1,00	0,01
100	0,897	0,95	0,32
250	0,524	0,72	0,69
500	0,079	0,28	0,96
1000	0,00004	0,006	1,00

К середине пути от картинки почти ничего не остаётся, а к шагу 1000 остаётся чистый шум. И главное удобство: зашумлённую картинку для любого $t$ получаем за один шаг, без цикла по 1…t.

Обратный процесс: очистка от шума

При очистке нужно по текущему $x_t$ и $t$ с помощью модели предсказать значение $x_{t-1}$ в предыдущий момент. На практике это делается так: модель $p_\theta$ предсказывает шум, добавленный в прямом процессе при переходе от $x_{t-1}$ к $x_t$:

$x_{t-1} = \frac{1}{\sqrt{\alpha_t}}\Big(x_t - \frac{1-\alpha_t}{\sqrt{1-\bar{\alpha}_t}}\,\epsilon_\theta(x_t,t)\Big) + \sigma_t z, \quad z \sim \mathcal{N}(0,I)$

где $\sigma_t^2 = \frac{(1-\alpha_t)(1-\bar{\alpha}_{t-1})}{1-\bar{\alpha}_t}$, а $\epsilon_\theta(x_t,t)$ — шум прямого процесса, который модель предсказывает по $x_t$ и $t$.

Это формулы (6)–(7) статьи (если подставить $x_t = \sqrt{\bar{\alpha}_t}x_0 + \sqrt{1-\bar{\alpha}_t}\,\epsilon_\theta(x_t,t)$ и воспользоваться $\bar{\alpha}_t = \prod_{i=1}^t \alpha_i$):

$q(x_{t-1} \mid x_t, x_0) = \mathcal{N}\big(x_{t-1};\ \tilde{\mu}_t(x_t, x_0),\ \tilde{\beta}_t I\big)$

где $\tilde{\mu}_t(x_t, x_0) := \frac{\sqrt{\bar{\alpha}_{t-1}}\,\beta_t}{1 - \bar{\alpha}_t}\,x_0 + \frac{\sqrt{\alpha_t}\,(1 - \bar{\alpha}_{t-1})}{1 - \bar{\alpha}_t}\,x_t$, $\tilde{\beta}_t := \frac{1 - \bar{\alpha}_{t-1}}{1 - \bar{\alpha}_t}\,\beta_t$.

Итог

Хотя вывод DDPM довольно сложен, на деле, чтобы при инференсе получить изображение, нужна только одна формула — формула для $x_{t-1}$ из обратного процесса. Её можно упрощённо записать так:

$x_{t-1} = a\,(x_t - b\,\epsilon_\theta(x_t,t)) + c\,z, \quad z \sim \mathcal{N}(0,I)$

где $a$, $b$, $c$ — известные, вычислимые константы, $\epsilon_\theta(x_t,t)$ — шум, предсказанный моделью, $z$ — шум из стандартного нормального распределения.

Цель обучения — чтобы предсказанный моделью шум $\epsilon_\theta(x_t,t)$ был как можно ближе к настоящему. Поэтому достаточно минимизировать среднеквадратичную ошибку (MSE) между предсказанным и настоящим шумом.

Примечание переводчика

Записанная функция потерь — это «упрощённая» цель из статьи DDPM: $L = \mathbb{E}_{t, x_0, \epsilon}\,\lVert \epsilon - \epsilon_\theta(\sqrt{\bar\alpha_t}x_0 + \sqrt{1-\bar\alpha_t}\epsilon,\ t)\rVert^2$. Интуиция, вывод и связь со score matching — в уроке 9.1.2 «Диффузия»; современная альтернатива с прямыми траекториями — 9.1.4 «Flow matching».

Реализация

1. Загрузка данных

Обучаем на датасете CIFAR-10. Сначала скачайте данные в папку datasets/ так, чтобы получилось datasets/cifar-10-batches-py.

Затем определяем преобразования данных. Для обучающей выборки — случайное отражение по горизонтали и приведение к размеру $32 \times 32$. Чтобы было удобно добавлять гауссов шум и убирать его, данные нужно перевести в диапазон $[-1,1]$. Для тестовой выборки — просто приводим к $32 \times 32$ и переводим в $[-1,1]$.

train_data_transform = transforms.Compose([
    transforms.Resize((img_size, img_size)),
    transforms.RandomHorizontalFlip(),  # случайное отражение по горизонтали
    transforms.ToTensor(),  # переводим данные в диапазон [0, 1]
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),  # переводим данные в диапазон [-1, 1]
])
test_data_transform = transforms.Compose([
    transforms.Resize((img_size, img_size)),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
])

Затем загружаем данные через torchvision.datasets.CIFAR10 и создаём DataLoader — подробности в ddpm/dataloader.py.

После этого можно написать функцию визуализации. Преобразованные данные переводим из диапазона $[-1,1]$ в $[0,255]$, меняем порядок каналов и превращаем в картинку PIL; если на входе батч — берём первую картинку.

def show_tensor_image(image):
    reverse_transforms = transforms.Compose([
        transforms.Lambda(lambda t: (t + 1) / 2),  # из [-1, 1] в [0, 1]
        transforms.Lambda(lambda t: t.permute(1, 2, 0)),  # порядок каналов CHW → HWC
        transforms.Lambda(lambda t: t * 255.),  # в диапазон [0, 255]
        transforms.Lambda(lambda t: t.numpy().astype(np.uint8)),  # в тип uint8
        transforms.ToPILImage(),  # в картинку PIL
    ])

    # Если на входе батч — берём первую картинку
    if len(image.shape) == 4:
        image = image[0, :, :, :]
    return reverse_transforms(image)

Проверить, что данные загружаются правильно, можно командой python ddpm/dataloader.py.

Примечание переводчика

Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) вычисляет $(x - 0{,}5)/0{,}5 = 2x - 1$: 0 → −1, 0,5 → 0, 1 → 1. Диапазон $[-1, 1]$ нужен, чтобы картинка и стандартный гауссов шум были «одного масштаба».

2. Модель

На вход модели DDPM подаются зашумлённое изображение $x_t$ и шаг времени $t$, на выходе — предсказанный шум $\epsilon_\theta(x_t,t)$.

Сначала определим слой временного эмбеддинга: он вносит информацию о времени в признаки, отображая шаг $t$ в вектор большой размерности. По аналогии с позиционным кодированием в Transformer используем синусы и косинусы:

$PE(t, 2i) = \sin(t / 10000^{2i/d})$ $PE(t, 2i+1) = \cos(t / 10000^{2i/d})$

где $d$ — размерность эмбеддинга, $i$ — номер координаты.

class SinusoidalPositionEmbeddings(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.dim = dim

    def forward(self, time):
        device = time.device
        # делим размерность пополам: одна половина для sin, другая для cos
        half_dim = self.dim // 2
        # шаг экспоненциального убывания частот
        embeddings = math.log(10000) / (half_dim - 1)
        # последовательность частот
        embeddings = torch.exp(torch.arange(half_dim, device=device) * -embeddings)
        # умножаем шаг времени на частоты
        embeddings = time[:, None] * embeddings[None, :]
        # склеиваем sin и cos — получаем итоговый эмбеддинг
        embeddings = torch.cat((embeddings.sin(), embeddings.cos()), dim=-1)
        return embeddings

Примечание переводчика

Это те же синусоидальные позиционные кодировки, что в уроке 4.5.2, только «позиция» здесь — номер шага шума. В коде синусы и косинусы идут двумя половинами, а не чередуются, как в формуле, — для сети это неважно.

Затем определим базовый модуль U-Net — Block — с временным эмбеддингом и повышением/понижением разрешения. Первая свёртка увеличивает число каналов, затем добавляется временной эмбеддинг, затем вторая свёртка смешивает признаки, и в конце — повышение или понижение разрешения. Мы используем упрощённую U-Net, без механизма внимания, который есть в модели исходной статьи.

class Block(nn.Module):
    def __init__(self, in_channels, out_channels, time_emb_dim, up=False):
        super().__init__()
        self.time_mlp = nn.Linear(time_emb_dim, out_channels)
        if up:
            self.conv1 = nn.Conv2d(2 * in_channels, out_channels, kernel_size=3, padding=1)  # из-за остаточных связей U-Net при повышении разрешения приклеиваются (concat) прежние признаки, поэтому входных каналов вдвое больше
            self.transform = nn.ConvTranspose2d(out_channels, out_channels, kernel_size=4, stride=2, padding=1)
        else:
            self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)
            self.transform = nn.Conv2d(out_channels, out_channels, kernel_size=4, stride=2, padding=1)

        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1)
        self.bnorm1 = nn.BatchNorm2d(out_channels)
        self.bnorm2 = nn.BatchNorm2d(out_channels)
        self.relu = nn.ReLU()

    def forward(self, x, t):
        # первая свёртка
        h = self.bnorm1(self.relu(self.conv1(x)))
        # временной эмбеддинг
        time_emb = self.relu(self.time_mlp(t))
        # вносим информацию о времени в карту признаков
        h = h + time_emb[..., None, None]
        # вторая свёртка
        h = self.bnorm2(self.relu(self.conv2(h)))
        # повышение или понижение разрешения
        return self.transform(h)

Наконец, собираем из нескольких Block модель U-Net. Сначала модель понижает разрешение, меняя число каналов 3 → 64 → 128 → 256 → 512 → 1024: каналов становится больше, карты признаков — меньше. Затем повышает разрешение с каналами 1024 → 512 → 256 → 128 → 64 → 3 (это выходные каналы; входных из-за остаточных связей вдвое больше): каналов становится меньше, карты признаков — больше. В каждый слой вносится информация о шаге времени, а на выходе получается предсказанный шум того же размера, что и входное изображение.

class SimpleUnet(nn.Module):
    def __init__(self):
        super().__init__()
        image_channels = 3
        down_channels = (64, 128, 256, 512, 1024)
        up_channels = (1024, 512, 256, 128, 64)
        out_dim = 3
        time_emb_dim = 32

        # слой временного эмбеддинга
        self.time_embed = nn.Sequential(
            SinusoidalPositionEmbeddings(time_emb_dim),
            nn.Linear(time_emb_dim, time_emb_dim),
        )

        # входной слой, слои понижения и повышения разрешения, выходной слой
        self.input = nn.Conv2d(image_channels, down_channels[0], kernel_size=3, padding=1)
        self.downs = nn.ModuleList([Block(down_channels[i], down_channels[i + 1], time_emb_dim) for i in range(len(down_channels) - 1)])
        self.ups = nn.ModuleList([Block(up_channels[i], up_channels[i + 1], time_emb_dim, up=True) for i in range(len(up_channels) - 1)])
        self.output = nn.Conv2d(up_channels[-1], out_dim, kernel_size=3, padding=1)

    def forward(self, x, time_step):
        # эмбеддинг шага времени
        t = self.time_embed(time_step)
        # начальная свёртка
        x = self.input(x)
        # прямой проход U-Net: сначала понижаем разрешение и собираем признаки, затем повышаем разрешение
        residual_stack = []
        for down in self.downs:
            x = down(x, t)
            residual_stack.append(x)
        for up in self.ups:
            residual_x = residual_stack.pop()
            x = torch.cat((x, residual_x), dim=1)
            x = up(x, t)
        return self.output(x)

Командой python ddpm/unet.py можно напечатать формы входа и выхода каждого слоя U-Net и проверить, что модель собрана правильно:

Input shape: torch.Size([1, 3, 32, 32])
Time embedding shape: torch.Size([1, 32])
After input conv shape: torch.Size([1, 64, 32, 32])

Downsampling process:
Down block 1 output shape: torch.Size([1, 128, 16, 16])
Down block 2 output shape: torch.Size([1, 256, 8, 8])
Down block 3 output shape: torch.Size([1, 512, 4, 4])
Down block 4 output shape: torch.Size([1, 1024, 2, 2])

Upsampling process:
Concatenated input shape before up block 1: torch.Size([1, 2048, 2, 2])
Up block 1 output shape: torch.Size([1, 512, 4, 4])
Concatenated input shape before up block 2: torch.Size([1, 1024, 4, 4])
Up block 2 output shape: torch.Size([1, 256, 8, 8])
Concatenated input shape before up block 3: torch.Size([1, 512, 8, 8])
Up block 3 output shape: torch.Size([1, 128, 16, 16])
Concatenated input shape before up block 4: torch.Size([1, 256, 16, 16])
Up block 4 output shape: torch.Size([1, 64, 32, 32])

Final output shape: torch.Size([1, 3, 32, 32])

flowchart LR
I["x_t<br/>3×32×32"] --> C0["conv 64×32×32"]
C0 --> D1["down 128×16×16"] --> D2["down 256×8×8"] --> D3["down 512×4×4"] --> D4["down 1024×2×2"]
D4 --> U1["up 512×4×4"] --> U2["up 256×8×8"] --> U3["up 128×16×16"] --> U4["up 64×32×32"] --> O["ε̂: 3×32×32"]
D4 -.->|concat| U1
D3 -.->|concat| U2
D2 -.->|concat| U3
D1 -.->|concat| U4
T["шаг t → эмбеддинг 32"] --> D1
T --> U1

Схема переводчика по выводу unet.py: пунктир — пропускающие связи U-Net (время подаётся во все блоки)

3. Обучение

Сначала нужно определить планировщик шума (noise scheduler): он управляет зашумлением и даёт зашумлённые изображения для разных шагов. По формулам выше реализуем его в коде. Для прямого процесса нужны $\beta_t$, $\alpha_t$, $\bar{\alpha}_t$, $\sqrt{\bar{\alpha}_t}$, $\sqrt{1-\bar{\alpha}_t}$. Задаём их через register_buffer — тогда они автоматически сохраняются и загружаются вместе с параметрами модели.

class NoiseScheduler(nn.Module):
    def __init__(self, beta_start=0.0001, beta_end=0.02, num_steps=1000):
        super().__init__()
        self.beta_start = beta_start
        self.beta_end = beta_end
        self.num_steps = num_steps

        # β_t: линейное расписание шума
        self.register_buffer('betas', torch.linspace(beta_start, beta_end, num_steps))
        # α_t = 1 - β_t
        self.register_buffer('alphas', 1.0 - self.betas)
        # α_bar_t = ∏(1-β_i) для i от 1 до t
        self.register_buffer('alpha_bar', torch.cumprod(self.alphas, dim=0))
        # sqrt(α_bar_t)
        self.register_buffer('sqrt_alpha_bar', torch.sqrt(self.alpha_bar))
        # sqrt(1-α_bar_t)
        self.register_buffer('sqrt_one_minus_alpha_bar', torch.sqrt(1.0 - self.alpha_bar))

Мы обучаем на батчах изображений, и эти величины нужно комбинировать с тензором картинок, а сейчас все они — одномерные тензоры. Поэтому их форму надо подогнать под тензоры другой размерности. Для этого в классе NoiseScheduler есть метод get: он берёт значение величины для заданных шагов и меняет форму; var — тензор величины, t — шаги, x_shape — целевая форма.

def get(self, var, t, x_shape):
    # берём значения величины для заданных шагов
    out = var[t]
    # меняем форму на [batch_size, 1, 1, 1] для трансляции (broadcasting)
    return out.view([t.shape[0]] + [1] * (len(x_shape) - 1))

Теперь можно реализовать зашумление по формуле $x_t = \sqrt{\bar{\alpha}_t}x_0 + \sqrt{1-\bar{\alpha}_t}\epsilon$, $\epsilon \sim \mathcal{N}(0,I)$, где $\sqrt{\bar{\alpha}_t}$ и $\sqrt{1-\bar{\alpha}_t}$ берутся методом get.

def add_noise(self, x, t):
    # sqrt(α_bar_t) для шага t
    sqrt_alpha_bar = self.get(self.sqrt_alpha_bar, t, x.shape)
    # sqrt(1-α_bar_t) для шага t
    sqrt_one_minus_alpha_bar = self.get(self.sqrt_one_minus_alpha_bar, t, x.shape)
    # шум из стандартного нормального распределения ε ~ N(0,I)
    noise = torch.randn_like(x)
    # прямой процесс диффузии: x_t = sqrt(α_bar_t) * x_0 + sqrt(1-α_bar_t) * ε
    return sqrt_alpha_bar * x + sqrt_one_minus_alpha_bar * noise, noise

Командой python ddpm/diffusion.py можно нарисовать, как изображение шаг за шагом зашумляется.

Рисунок в оригинале (fig/fig2.png): одна и та же картинка CIFAR-10 на нескольких шагах $t$ — от чёткой до сплошного шума.

Наконец, можно написать весь цикл обучения. Шаги такие:

1. Случайно выбрать шаг времени t.
2. Зашумить изображения — получить зашумлённые картинки и сам шум.
3. Предсказать шум моделью.
4. Посчитать MSE между предсказанным и настоящим шумом.
5. Обратное распространение и шаг оптимизатора.

Командой python ddpm/train.py модель обучается и сохраняет параметры; 200 эпох на RTX-4090 занимают около 2 часов.

Примечание переводчика: шаг обучения в коде

Сам train.py в README не приведён; пять шагов выше выглядят примерно так (наш набросок по описанию):

for x0, _ in dataloader:                       # метки классов CIFAR-10 не нужны
    x0 = x0.to(device)
    t = torch.randint(0, scheduler.num_steps, (x0.shape[0],), device=device)  # 1. случайный шаг для каждой картинки
    x_t, noise = scheduler.add_noise(x0, t)    # 2. зашумили, запомнили шум
    pred = model(x_t, t)                       # 3. модель угадывает шум
    loss = F.mse_loss(pred, noise)             # 4. MSE с настоящим шумом
    optimizer.zero_grad(); loss.backward(); optimizer.step()  # 5. шаг оптимизатора

Это та же «петля обучения» из урока 3.3.2 — необычна только постановка задачи.

4. Генерация (сэмплирование)

Идея генерации: берём начальный шум из стандартного нормального распределения, затем шаг за шагом убираем шум от $t=T$ до $t=0$ и в конце обрезаем результат до диапазона $[-1,1]$. Для очистки нужны коэффициенты $\sqrt{\frac{1}{\bar{\alpha}_t}}$, $\sqrt{\frac{1}{\bar{\alpha}_t}-1}$, $\mu_\theta(x_t,t)$, $\log(\sigma_t^2)$. Определяем их в классе NoiseScheduler:

# α_bar_(t-1)
self.register_buffer('alpha_bar_prev', torch.cat([torch.tensor([1.0]), self.alpha_bar[:-1]], dim=0))
# 1/sqrt(α_t)
self.register_buffer('sqrt_recip_alphas', torch.sqrt(1.0 / self.alphas))

# 1/sqrt(α_bar_t)
self.register_buffer('sqrt_recip_alphas_bar', torch.sqrt(1.0 / self.alpha_bar))
# sqrt(1/α_bar_t - 1)
self.register_buffer('sqrt_recipm1_alphas_bar', torch.sqrt(1.0 / self.alpha_bar - 1))
# дисперсия апостериорного распределения σ_t^2
self.register_buffer('posterior_var', self.betas * (1.0 - self.alpha_bar_prev) / (1.0 - self.alpha_bar))
# коэффициент 1 среднего апостериорного распределения: β_t * sqrt(α_bar_(t-1))/(1-α_bar_t)
self.register_buffer('posterior_mean_coef1', self.betas * torch.sqrt(self.alpha_bar_prev) / (1.0 - self.alpha_bar))
# коэффициент 2 среднего апостериорного распределения: (1-α_bar_(t-1)) * sqrt(α_t)/(1-α_bar_t)
self.register_buffer('posterior_mean_coef2', (1.0 - self.alpha_bar_prev) * torch.sqrt(self.alphas) / (1.0 - self.alpha_bar))

Порядок генерации (по формулам статьи):

1. Взять начальный шум из стандартного нормального распределения: $x_T \sim \mathcal{N}(0,I)$.
2. Повторять в цикле от $t=T$ до $t=0$ следующие шаги:
3. По текущему шагу $t$ и текущему изображению $x_t$ вычислить моделью предсказанный шум $\epsilon_\theta(x_t,t)$.
4. Вычислить предсказание $x_0$: $x_0 = \frac{1}{\sqrt{\bar{\alpha}_t}}x_t - \sqrt{\frac{1}{\bar{\alpha}_t}-1}\,\epsilon_\theta(x_t,t)$.
5. Вычислить среднее апостериорного распределения: $\mu_\theta(x_t,t) = \frac{\sqrt{\bar{\alpha}_{t-1}}\beta_t}{1-\bar{\alpha}_t}x_0 + \frac{\sqrt{\alpha_t}(1-\bar{\alpha}_{t-1})}{1-\bar{\alpha}_t}x_t$.
6. Вычислить логарифм дисперсии апостериорного распределения: $\log(\sigma_t^2) = \log(\tilde{\beta}_t) = \log\!\Big(\frac{\beta_t(1-\bar{\alpha}_{t-1})}{1-\bar{\alpha}_t}\Big)$.
7. Если $t>0$, сделать выборку из апостериорного распределения: $x_{t-1} = \mu_\theta(x_t,t) + \sigma_t\epsilon$, $\epsilon \sim \mathcal{N}(0,I)$.
8. Если $t=0$, взять среднее как результат генерации: $x_0 = \mu_\theta(x_t,t)$.

def sample(model, scheduler, num_samples, size, device="cpu"):
    model.eval()
    with torch.no_grad():
        # начальный шум из стандартного нормального распределения x_T ~ N(0,I)
        x_t = torch.randn(num_samples, *size).to(device)

        # шаг за шагом убираем шум, от t=T до t=0
        for t in tqdm(reversed(range(scheduler.num_steps)), desc="Sampling"):
            # батч из номеров шага
            t_batch = torch.tensor([t] * num_samples).to(device)

            # коэффициенты для генерации
            sqrt_recip_alpha_bar = scheduler.get(scheduler.sqrt_recip_alphas_bar, t_batch, x_t.shape)
            sqrt_recipm1_alpha_bar = scheduler.get(scheduler.sqrt_recipm1_alphas_bar, t_batch, x_t.shape)
            posterior_mean_coef1 = scheduler.get(scheduler.posterior_mean_coef1, t_batch, x_t.shape)
            posterior_mean_coef2 = scheduler.get(scheduler.posterior_mean_coef2, t_batch, x_t.shape)

            # предсказываем шум ε_θ(x_t,t)
            predicted_noise = model(x_t, t_batch)

            # предсказание x_0: x_0 = 1/sqrt(α_bar_t) * x_t - sqrt(1/α_bar_t-1) * ε_θ(x_t,t)
            _x_0 = sqrt_recip_alpha_bar * x_t - sqrt_recipm1_alpha_bar * predicted_noise
            # среднее апостериорного распределения μ_θ(x_t,t)
            model_mean = posterior_mean_coef1 * _x_0 + posterior_mean_coef2 * x_t
            # логарифм дисперсии log(σ_t^2)
            model_log_var = scheduler.get(torch.log(torch.cat([scheduler.posterior_var[1:2], scheduler.betas[1:]])), t_batch, x_t.shape)

            if t > 0:
                # при t>0 делаем выборку из апостериорного распределения: x_t-1 = μ_θ(x_t,t) + σ_t * z, z~N(0,I)
                noise = torch.randn_like(x_t).to(device)
                x_t = model_mean + torch.exp(0.5 * model_log_var) * noise
            else:
                # при t=0 берём среднее как результат
                x_t = model_mean
        # обрезаем результат до диапазона [-1,1]
        x_0 = torch.clamp(x_t, -1.0, 1.0)
    return x_0

Примечание переводчика: какая дисперсия на самом деле

Обратите внимание на строку model_log_var = ... torch.cat([posterior_var[1:2], betas[1:]]). В тексте сказано $\sigma_t^2 = \tilde\beta_t$, а код берёт $\sigma_t^2 = \beta_t$ (кроме самого первого шага). Это не ошибка, а второй вариант из статьи DDPM: авторы статьи сообщают, что $\sigma_t^2 = \beta_t$ и $\sigma_t^2 = \tilde\beta_t$ дают похожие результаты. Так же сделано в исходном коде DDPM (вариант fixedlarge).

И ещё: генерация — это 1000 последовательных проходов U-Net на каждую партию картинок, поэтому она медленная. Методы ускорения (DDIM, меньше шагов) — в уроке 9.1.2.

Командой python ddpm/sample.py можно загрузить обученную модель, сгенерировать изображения и сохранить их.

Рисунок в оригинале (fig/fig3.png): сетка сгенерированных картинок 32×32.

5. Оценка

Inception Score (IS) и Fréchet Inception Distance (FID) — две важные метрики качества сгенерированных изображений; по ним мы и оцениваем качество генерации.

Сначала загружаем предобученную модель Inception и получаем признаки настоящих и сгенерированных изображений.

class InceptionStatistics:
    def __init__(self, device='cuda'):
        self.device = device
        # загружаем предобученную модель Inception v3
        self.model = models.inception_v3(weights=models.Inception_V3_Weights.IMAGENET1K_V1, transform_input=False)
        self.model.fc = nn.Identity()  # убираем последний полносвязный слой
        self.model = self.model.to(device)
        self.model.eval()

        # предобработка изображений
        self.preprocess = transforms.Compose([
            transforms.Resize(299),
            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
        ])

    @torch.no_grad()
    def get_features(self, images):
        """Получить признаки Inception"""
        features = []
        probs = []

        # приводим изображения к размеру 299x299
        images = self.preprocess(images)

        # обрабатываем изображения батчами
        dataset = TensorDataset(images)
        dataloader = DataLoader(dataset, batch_size=32)

        for (batch,) in tqdm(dataloader):
            batch = batch.to(self.device)

            # получаем признаки и логиты
            feature = self.model(batch)
            prob = F.softmax(feature, dim=1)

            features.append(feature.cpu().numpy())
            probs.append(prob.cpu().numpy())

        features = np.concatenate(features, axis=0)
        probs = np.concatenate(probs, axis=0)

        return features, probs

Inception Score (IS)

IS оценивает качество и разнообразие сгенерированных изображений с помощью предобученной сети Inception v3. Формула:

$IS = \exp\big(\mathbb{E}_{x\sim p_g}[\mathrm{KL}(p(y\mid x) \,\Vert\, p(y))]\big)$

где:

• $p_g$ — распределение генератора;
• $p(y\mid x)$ — вероятности классов, которые модель Inception предсказывает для изображения $x$;
• $p(y)$ — среднее распределение классов по всем сгенерированным изображениям;
• KL — расхождение Кульбака — Лейблера, мера различия двух распределений вероятностей: $\mathrm{KL}(p \,\Vert\, q) = \sum_{i=1}^{n} p(i) \log \frac{p(i)}{q(i)}$.

Чем выше IS, тем:

1. увереннее предсказание класса для каждого изображения (качество выше);
2. шире разброс классов между изображениями (разнообразие выше).

Шаги:

1. Разбить все изображения на группы.
2. Для каждой группы:
   • посчитать маргинальное распределение $p(y)$ — среднее $p(y\mid x)$ по группе;
   • посчитать KL-расхождение;
   • взять экспоненту.
3. Вернуть среднее и стандартное отклонение оценок по группам.

def calculate_inception_score(probs, splits=10):
    # оценки IS для каждой группы
    scores = []
    # размер группы
    split_size = probs.shape[0] // splits

    # считаем для каждой группы
    for i in tqdm(range(splits)):
        # распределения вероятностей текущей группы
        part = probs[i * split_size:(i + 1) * split_size]
        # KL-расхождение: KL(p(y|x) || p(y))
        kl = part * (np.log(part) - np.log(np.expand_dims(np.mean(part, axis=0), 0)))
        # усредняем KL по объектам
        kl = np.mean(np.sum(kl, axis=1))
        # exp(KL) — в список оценок
        scores.append(np.exp(kl))

    # среднее и стандартное отклонение IS по группам
    return np.mean(scores), np.std(scores)

Fréchet Inception Distance (FID)

FID оценивает качество генерации, сравнивая распределения настоящих и сгенерированных изображений в пространстве признаков Inception. Формула:

$FID = \lVert\mu_r - \mu_g\rVert^2 + \mathrm{Tr}\big(\Sigma_r + \Sigma_g - 2(\Sigma_r\Sigma_g)^{1/2}\big)$

где:

• $\mu_r, \mu_g$ — средние признаков настоящих и сгенерированных изображений;
• $\Sigma_r, \Sigma_g$ — ковариационные матрицы признаков настоящих и сгенерированных изображений;
• $\mathrm{Tr}$ — след матрицы.

Чем ниже FID, тем ближе распределение признаков сгенерированных изображений к настоящему и тем выше качество.

Шаги:

1. Для настоящих и для сгенерированных изображений по отдельности:
   • извлечь признаки моделью Inception;
   • посчитать вектор средних и ковариационную матрицу признаков.
2. Посчитать евклидово расстояние между векторами средних.
3. Посчитать квадратный корень из произведения ковариационных матриц.
4. Посчитать итоговый FID.

def calculate_fid(real_features, fake_features):
    # векторы средних и ковариационные матрицы признаков настоящих и сгенерированных изображений
    mu1, sigma1 = real_features.mean(axis=0), np.cov(real_features, rowvar=False)
    mu2, sigma2 = fake_features.mean(axis=0), np.cov(fake_features, rowvar=False)

    # квадрат евклидова расстояния между векторами средних
    ssdiff = np.sum((mu1 - mu2) ** 2)
    # квадратный корень из произведения ковариаций: (Σ_r Σ_f)^(1/2)
    covmean = linalg.sqrtm(sigma1.dot(sigma2))  # считается долго
    # если в результате есть комплексные числа — берём действительную часть
    if np.iscomplexobj(covmean):
        covmean = covmean.real

    # итоговый FID
    fid = ssdiff + np.trace(sigma1 + sigma2 - 2 * covmean)
    return fid

Командой python ddpm/metrics.py можно посчитать IS и FID: мы генерируем обученной диффузионной моделью 10 000 изображений и сравниваем их с датасетом CIFAR-10 как набором настоящих изображений.

По одной картинке IS и FID не посчитать, но можно наглядно посмотреть, как выглядят изображения двух моделей с разными IS и FID. Ниже — изображения упрощённой U-Net этого проекта (epochs=2000, IS=1.12, FID=41.63) и U-Net со слоями внимания (Attention) из статьи DDPM (epochs=2000, IS=1.17, FID=14.10), обученных в нашем фреймворке 200 и 2000 эпох.

Рисунки в оригинале — четыре сетки сгенерированных картинок: SimpleUnet, 200 эпох, AttentionUnet, 200 эпох, SimpleUnet, 2000 эпох, AttentionUnet, 2000 эпох.

Видно, что по мере обучения и с ростом качества модели изображения модели с более высоким IS и более низким FID лучше узнаются по классу, а детали ближе к настоящим картинкам CIFAR-10.

Примечание переводчика: почему IS здесь около 1

В InceptionStatistics последний слой fc заменён на nn.Identity(), поэтому feature — это 2048-мерные признаки перед классификатором, а не логиты 1000 классов ImageNet. Для FID так и нужно, а вот softmax от этих признаков — не вероятности классов. Поэтому IS у авторов получился 1,12–1,17 (минимум IS равен 1). Для настоящего IS нужны логиты исходного слоя fc; в литературе у DDPM на CIFAR-10 IS = 9,46 и FID = 3,17 (по 50 000 картинок). Сравнивать числа этого проекта можно только между собой: FID 41,6 против 14,1 честно показывает, что U-Net с вниманием лучше.

От переводчика: как запустить на своей машине

🖥 На вашей машине (4×GPU)

• Авторская оценка: 200 эпох на RTX 4090 — около 2 часов. На одной RTX 4080 ожидайте того же порядка (немного дольше). Три остальные карты можно отдать под параллельные опыты: другое расписание $\beta_t$, AttentionUnet, меньше каналов.
• Для быстрой проверки поставьте 5–10 эпох: картинки будут размытыми «пятнами», но цикл «обучение → генерация» заработает.
• metrics.py генерирует 10 000 картинок по 1000 шагов и считает sqrtm матрицы 2048×2048 — это самый долгий шаг; для пробы уменьшите число картинок до 1000.

Что сделать руками:

1. Повторите таблицу $\bar\alpha_t$ из примечания в NumPy и постройте график $\sqrt{\bar\alpha_t}$ по $t$ для линейного расписания. Сравните с косинусным расписанием (Improved DDPM).
2. Обучите SimpleUnet 20 эпох и сгенерируйте 16 картинок с 1000 шагами. Затем сократите генерацию до 100 шагов, просто пропуская шаги, — и посмотрите, почему так делать нельзя без DDIM.
3. Исправьте подсчёт IS (верните fc, а признаки для FID берите отдельно) и пересчитайте метрики.

Словарик

Русский	Пояснение
диффузионная модель (diffusion model)	генерация через постепенное удаление шума
прямой процесс, зашумление (forward process)	$x_0 \to x_T$
обратный процесс, очистка от шума (reverse process)	$x_T \to x_0$
марковская цепь (Markov chain)	следующее состояние зависит только от текущего
расписание шума (noise schedule)	значения $\beta_t$ по шагам
апостериорное распределение (posterior)	$q(x_{t-1}\mid x_t, x_0)$
временной эмбеддинг (time embedding)	номер шага → вектор
повышение / понижение разрешения	upsampling / downsampling
генерация, сэмплирование (sampling)	получение новых картинок
расхождение Кульбака — Лейблера (KL divergence)	урок 1.3.3