Содержание курса 31 / 33

31 / 33 · Авторегрессионные модели

Авторегрессионные модели: PixelRNN и PixelCNN

Плотность изображения как произведение одномерных условных распределений, каждое из которых — softmax по значениям пикселя: точное правдоподобие, биты на размерность, параллельное обучение и последовательное сэмплирование, рекуррентные (PixelRNN) и маскированные свёрточные (PixelCNN) сети, слепое пятно, Gated PixelCNN и PixelCNN++.

На этой странице 12

Семейства этого курса по-разному обходятся с невычислимым правдоподобием: VAE его оценивает снизу, GAN избегает, поток вычисляет через обратимое отображение, диффузия снова оценивает снизу. Авторегрессионные модели убирают трудность у самого источника. Они записывают плотность изображения как произведение одномерных условных распределений, каждое из которых — softmax по нескольким сотням значений и потому нормировано по построению. Правдоподобие точное, обучение сводится к кросс-энтропии, а платить приходится при сэмплировании: пиксель за пикселем.

Обозначения. Изображение — вектор из nn значений x=(x1,…,xn)x = (x_1, \dots, x_n) в фиксированном порядке, обычно построчно от левого верхнего угла, а у цветных изображений ещё и красный, зелёный, синий внутри каждого пикселя. x<i=(x1,…,xi−1)x_{<i} = (x_1, \dots, x_{i-1}) — значения до позиции ii. Значения — целые числа из {0,…,255}\{0, \dots, 255\}, если не сказано иное.

Цепное правило как модель#

Любое совместное распределение можно записать как

p(x)=∏i=1np(xi∣x<i),log⁡p(x)=∑i=1nlog⁡p(xi∣x<i).p(x) = \prod_{i=1}^{n} p(x_i \mid x_{<i}), \qquad \log p(x) = \sum_{i=1}^{n} \log p(x_i \mid x_{<i}).

Это цепное правило теории вероятностей, а не допущение: оно верно для любого распределения и любого порядка координат. Ограничения модели появляются только из того, как параметризованы условные распределения. Метод максимального правдоподобия, который заметка о KL выводит из минимизации DKL(pdata ∥ pθ)D_{\mathrm{KL}}(p_{\text{data}} \,\|\, p_\theta), превращается в сумму nn задач классификации: в каждой позиции предсказать следующее значение по предыдущим и заплатить кросс-энтропию истинного значения.

Вся трудность — в компактном представлении условных распределений. Возьмём бинарные изображения из nn пикселей. Таблица совместного распределения содержит 2n−12^n - 1 свободных чисел, и столько же в сумме содержат таблицы всех условных распределений: условному распределению пикселя ii нужно по числу на каждую из 2i−12^{i-1} предысторий. Полностью видимая сигмоидная сеть доверия (fully visible sigmoid belief network) заменяет каждую таблицу логистической регрессией,

p(xi=1∣x<i)=σ(αi+∑j<iαijxj),p(x_i = 1 \mid x_{<i}) = \sigma\Big(\alpha_i + \sum_{j<i} \alpha_{ij} x_j\Big),

с ii параметрами для пикселя ii и n(n+1)/2n(n+1)/2 параметрами всего. Для изображений MNIST размера 28×2828 \times 28, n=784n = 784, это 307 720 параметров вместо 2784−12^{784} - 1. Более поздние модели разделяют параметры сильнее: NADE (Larochelle and Murray, 2011) вычисляет все условные распределения одним скрытым слоем, веса которого переиспользуются между позициями, а MADE (Germain et al., 2015) получает их все за один проход автоэнкодера с масками на весах, так что выход ii зависит только от входов до ii; те же маски строят авторегрессионные потоки. PixelRNN и PixelCNN идут дальше и применяют одну рекуррентную или свёрточную сеть во всех позициях, так что число параметров перестаёт зависеть от размера изображения.

Сэмплирование следует произведению: вытянуть x1x_1, затем x2x_2 при данном x1x_1, и так далее. Каждый шаг — метод обращения функции распределения для дискретной величины: взять u∼U(0,1)u \sim \mathcal U(0, 1) и вернуть наименьшее значение, накопленная вероятность которого достигает uu. Авторегрессионная модель — по-прежнему шум плюс обучаемое преобразование, где шумом служат nn равномерных чисел.

Пиксели как дискретные значения#

PixelRNN (van den Oord, Kalchbrenner and Kavukcuoglu, 2016) считает каждое значение цвета категорией и предсказывает его softmax на 256 классов. Дискретному распределению не нужны допущения о форме: оно может быть многомодальным или скошенным, и оно назначает вероятности, а не плотности: log⁡p(x)≤0\log p(x) \le 0, и деквантование (dequantization) не нужно, в отличие от непрерывных плотностей потоков.

Правдоподобие моделей изображений сообщают в битах на размерность (bits/dim):

bits/dim=−log⁡2p(x)n=−ln⁡p(x)nln⁡2,\text{bits/dim} = -\frac{\log_2 p(x)}{n} = -\frac{\ln p(x)}{n \ln 2},

с усреднением по тестовым изображениям. Модель, распределяющая вероятность равномерно по 256 значениям, получает ровно 8 бит на размерность. У этого числа прямой смысл: с арифметическим кодированием модель с bb битами на размерность сжимает изображения без потерь примерно до b⋅nb \cdot n бит. Для CIFAR-10 n=32⋅32⋅3=3072n = 32 \cdot 32 \cdot 3 = 3072, и лучший PixelRNN достигает 3.00 бит на размерность: 3.00⋅3072=92163.00 \cdot 3072 = 9216 бит, около 1.1 КиБ на изображение, или 37.5% от 24 576 бит исходного изображения.

Модель, предсказывающая меньше уровней, решает более грубую задачу. При 10 уровнях серого равномерное угадывание стоит ln⁡10≈2.30\ln 10 \approx 2.30 ната, или log⁡210≈3.32\log_2 10 \approx 3.32 бита на пиксель, и сравнивать такие числа с битами на размерность по 256 значениям нельзя.

Параллельное обучение, последовательное сэмплирование#

При обучении изображение известно целиком. Если сеть устроена так, что её выход в позиции ii зависит только от входов x<ix_{<i}, один прямой проход по настоящему изображению даёт сразу все nn условных распределений, а функция потерь — их сумма. Подача на вход истинных предыдущих значений, а не собственных сэмплов модели, называется teacher forcing, и благодаря ей обучение так же параллельно, как у любой свёрточной сети.

Сэмплирование параллельным быть не может: xi+1x_{i+1} вычисляется по xix_i, который нужно сначала вытянуть. Модели, которая предсказывает по одному значению, на сэмпл нужно nn проходов сети: 3072 для изображения CIFAR-10 и 89⋅60=534089 \cdot 60 = 5340 для серого изображения 89×6089 \times 60 пикселей. Модели, которая за проход предсказывает целый пиксель, как PixelCNN++ ниже, нужен один проход на пиксель, 1024 для CIFAR-10, а три значения пикселя она вытягивает из предсказанных параметров; позиции по-прежнему идут последовательно. Простая реализация к тому же на каждом проходе пересчитывает всё изображение, хотя использует один выход. Та же асимметрия — дешёвая плотность и дорогое сэмплирование — есть у маскированного авторегрессионного потока.

PixelRNN и PixelCNN#

Две сети из van den Oord, Kalchbrenner and Kavukcuoglu (2016) различаются тем, как выход в позиции ii собирает информацию из всего x<ix_{<i}, ничего не видя из xix_i и последующих значений.

Рекуррентные слои#

Row LSTM обрабатывает изображение строка за строкой сверху вниз и вычисляет состояние сразу для всей строки. Переход вход–состояние — свёртка k×1k \times 1 вдоль строки (k≥3k \ge 3), а переход состояние–состояние — одномерная свёртка по предыдущей строке состояний. Каждый шаг смотрит немного влево и вправо в строке выше, поэтому контекст пикселя — треугольник над ним, и часть доступного прошлого в него не попадает.

Diagonal BiLSTM охватывает всё прошлое. Он скашивает входную карту, сдвигая каждую строку на одну позицию вправо относительно предыдущей, так что диагонали изображения становятся столбцами карты n×(2n−1)n \times (2n-1). Свёртка состояние–состояние имеет размер 2×12 \times 1 вдоль этих столбцов, и два LSTM сканируют в противоположных направлениях; выход прохода справа налево сдвигается на строку вниз перед сложением, чтобы не просочились будущие пиксели. В моделях до 12 слоёв с остаточными связями между ними. Рекуррентности идут последовательно по строкам или диагоналям даже при обучении, поэтому обе модели обучаются медленно.

Маскированные свёртки#

PixelCNN заменяет рекуррентности свёртками, ядра которых умножаются на маску. Для ядра k×kk \times k с центром c=⌊k/2⌋c = \lfloor k/2 \rfloor

Muv={1,u<c или (u=c и v<c),0,иначе,M_{uv} = \begin{cases} 1, & u < c \text{ или } (u = c \text{ и } v < c), \\ 0, & \text{иначе}, \end{cases}

так что позиция видит строки над собой и пиксели слева от себя. Первый слой использует маску A, которая исключает центр: на входе первого слоя в центре стоит значение, которое нужно предсказать. Последующие слои используют маску B, которая центр включает: после первого слоя признак в позиции ii вычислен только по x<ix_{<i}, и передавать его дальше безопасно. PixelCNN из статьи состоит из пятнадцати слоёв: первый 7×77 \times 7 с маской A, затем остаточные блоки со свёртками 3×33 \times 3. Обучение полностью параллельно, сэмплирование — нет.

Рисунок

Маски свёртки 5×5 в PixelCNN

Маски типа A и BДве сетки 5×5 из нулей и единиц. Единицы в двух верхних строках и в двух левых клетках средней строки; в маске B единица ещё и в центре.маска A: первый слой1111111111110000000000000маска B: остальные слои1111111111111000000000000центр: предсказываемый пиксель
Маска умножается на ядро свёртки поэлементно, и веса с нулями не участвуют в вычислении. Единицы стоят во всех строках выше центра и слева от центра в его строке: это пиксели, предшествующие текущему в порядке развёртки. Маска A исключает и сам центр (12 единиц из 25), потому что на входе первого слоя в центре стоит значение, которое нужно предсказать. Маска B центр оставляет (13 единиц): после первого слоя признак в позиции i уже вычислен только по пикселям до i.

Цветовые каналы#

В цветном изображении каждый пиксель содержит три значения в порядке R, G, B, и факторизация идёт по всем 3×3 \times (число пикселей) значениям:

p(xi∣x<i)=p(xi,R∣x<i)  p(xi,G∣x<i,xi,R)  p(xi,B∣x<i,xi,R,xi,G),p(x_i \mid x_{<i}) = p(x_{i,R} \mid x_{<i})\; p(x_{i,G} \mid x_{<i}, x_{i,R})\; p(x_{i,B} \mid x_{<i}, x_{i,R}, x_{i,G}),

где xix_i теперь весь пиксель. Карты признаков каждого слоя делятся на три группы, по одной на цвет. Маска A не даёт красной группе видеть ни одного значения текущего пикселя, зелёной позволяет видеть красный, а синей — красный и зелёный; маска B вдобавок соединяет каждую группу с самой собой.

Простейшая альтернатива моделирует каналы независимо, например тремя моделями полутоновых изображений, по одной на канал. Тогда каждый канал сам по себе — правдоподобное изображение, но ничто не связывает три канала между собой: в сэмпле красное, зелёное и синее значения пикселя, а также границы, которые они рисуют, не обязаны согласовываться. Более мягкий вариант ставит в условие каждого канала все предыдущие пиксели всех каналов, но отбрасывает зависимость между тремя значениями текущего пикселя. Цепному правилу эта зависимость нужна, поэтому такой вариант — допущение, а не факторизация. Модель, сохраняющая эту зависимость, либо соблюдает порядок каналов внутри масок, как описано выше, либо, как PixelCNN++ ниже, предсказывает три значения пикселя совместно с простой зависимостью между ними.

Слепое пятно#

Стопка маскированных свёрток не даёт всего прошлого. Маска 3×33 \times 3 дотягивается на строку вверх и на столбец в каждую сторону или на столбец влево в текущей строке. После LL слоёв строка на aa выше цели покрыта от LL столбцов слева лишь до aa столбцов справа, каким бы ни было LL. Пиксели правее в ближайших строках сверху остаются недосягаемыми при любой глубине. В окне из LL строк и ±L\pm L столбцов это слепое пятно содержит L(L−1)/2L(L-1)/2 из 2L(L+1)2L(L+1) предшествующих пикселей, 6 из 40 при L=4L = 4; его доля растёт к четверти — эту цифру приводят van den Oord et al. (2016b). Раздел с кодом ниже находит те же 6 пикселей, дифференцируя сеть со случайной инициализацией.

Рисунок

Слепое пятно PixelCNN

Рецептивное поле и слепое пятноДве сетки 6×11 с предсказываемым пикселем в нижней строке. В сетке PixelCNN часть прошлого сверху справа (6 пикселей) не входит в рецептивное поле. В сетке Gated PixelCNN рецептивное поле покрывает всё прошлое в пределах 4 строк.PixelCNN: 4 маскированные свёртки 3×3Gated PixelCNN: две стопки, 4 слоявидимая часть прошлогослепое пятнопредсказываемый пиксельРецептивное поле и слепое пятноДве сетки 6×11 с предсказываемым пикселем в нижней строке. В сетке PixelCNN часть прошлого сверху справа (6 пикселей) не входит в рецептивное поле. В сетке Gated PixelCNN рецептивное поле покрывает всё прошлое в пределах 4 строк.PixelCNN: 4 маскированные свёртки 3×3Gated PixelCNN: две стопки, 4 слоявидимая часть прошлогослепое пятнопредсказываемый пиксель
Какие предшествующие пиксели могут повлиять на предсказание в отмеченном пикселе после 4 слоёв. Слева — PixelCNN: в первой свёртке 3×3 маска типа A, в остальных — типа B. В строке на a выше цели поле доходит вправо лишь на a столбцов, поэтому треугольник сверху справа, хотя он целиком в прошлом, остаётся вне поля при любой глубине: в этом окне это 6 пикселей. Справа — Gated PixelCNN: вертикальная стопка покрывает строки выше, горизонтальная — пиксели слева в той же строке, и слепого пятна нет.

Gated PixelCNN#

van den Oord et al. (2016b) закрывают разрыв между PixelCNN и PixelRNN двумя изменениями. Первое устраняет слепое пятно двумя стопками свёрток. Вертикальная стопка не маскируется внутри своего окна и видит все строки выше текущей, так что её рецептивное поле растёт прямоугольником. Горизонтальная стопка видит текущую строку до текущего пикселя. После каждого слоя вертикальная стопка передаёт выход горизонтальной, но не наоборот: информация из горизонтальной стопки занесла бы в вертикальную пиксели правее или ниже цели.

Второе изменение заменяет ReLU между слоями мультипликативными блоками, аналогом гейтов LSTM:

y=tanh⁡(Wk,f∗x)⊙σ(Wk,g∗x),y = \tanh(W_{k,f} * x) \odot \sigma(W_{k,g} * x),

где ∗* — свёртка, kk — номер слоя, ⊙\odot — поэлементное произведение. Условие входит дополнительным слагаемым под обеими нелинейностями: для вектора hh, например one-hot метки класса, y=tanh⁡(Wk,f∗x+Vk,f⊤h)⊙σ(Wk,g∗x+Vk,g⊤h)y = \tanh(W_{k,f} * x + V_{k,f}^\top h) \odot \sigma(W_{k,g} * x + V_{k,g}^\top h), а для пространственной карты ss произведения V⊤hV^\top h становятся свёртками V∗sV * s. Условная модель — p(x∣h)=∏ip(xi∣x<i,h)p(x \mid h) = \prod_i p(x_i \mid x_{<i}, h).

На CIFAR-10 Gated PixelCNN достигает 3.03 бита на размерность против 3.14 у PixelCNN и 3.00 у PixelRNN. На ImageNet он немного лучше PixelRNN: 3.83 против 3.86 при 32×3232 \times 32 и 3.57 против 3.63 при 64×6464 \times 64, а качества Row LSTM достигает менее чем за половину времени обучения (60 часов на 32 GPU). Модель для ImageNet — 20 слоёв по 384 канала с фильтрами 5×55 \times 5.

PixelCNN++#

Softmax на 256 классов ничего не знает о порядке своих классов: 127 и 128 для него так же не связаны, как 0 и 255, пока сеть не выучит иное; значение, не встреченное при обучении, получает почти нулевую вероятность; а выходной слой содержит 3⋅256=7683 \cdot 256 = 768 чисел на пиксель. PixelCNN++ (Salimans et al., 2017) предполагает непрерывную интенсивность ν\nu со смесью KK логистических распределений и моделирует наблюдаемое значение как ν\nu, округлённое до целого:

P(x∣π,μ,s)=∑k=1Kπk[σ(x+0.5−μksk)−σ(x−0.5−μksk)],P(x \mid \pi, \mu, s) = \sum_{k=1}^{K} \pi_k \Big[\sigma\Big(\frac{x + 0.5 - \mu_k}{s_k}\Big) - \sigma\Big(\frac{x - 0.5 - \mu_k}{s_k}\Big)\Big],

где σ\sigma — логистическая сигмоида, функция распределения логистического закона. Для x=0x = 0 слагаемое x−0.5x - 0.5 заменяется на −∞-\infty, а для x=255x = 255 слагаемое x+0.5x + 0.5 — на +∞+\infty, так что вероятности по {0,…,255}\{0, \dots, 255\} дают в сумме единицу, а крайние значения собирают массу всего, что лежит за ними. С одной компонентой, μ=127.3\mu = 127.3 и s=2s = 2, значения 126, 127, 128 и 129 получают вероятности 0.112, 0.124, 0.121 и 0.105: соседи близки автоматически. При μ=254\mu = 254 и s=2s = 2 крайнее значение 255 получает 0.438, а 254 — 0.124, что соответствует избытку насыщенных пикселей в реальных фотографиях.

Остальные изменения:

  • Целые пиксели. Сеть обусловливается на целые пиксели сверху и слева и предсказывает три канала пикселя вместе. Индикатор компоненты смеси общий для каналов, а внутри компоненты средние линейно зависят от уже вытянутых каналов: μG+αR\mu_G + \alpha R и μB+βR+γG\mu_B + \beta R + \gamma G, где α,β,γ\alpha, \beta, \gamma выдаёт сеть. Выход — 10K10K чисел на пиксель (KK весов, среднее и масштаб для каждого канала и каждой компоненты и три коэффициента на компоненту), 50 при K=5K = 5, чего, по данным статьи, достаточно.
  • Понижение разрешения свёртками с шагом 2 вместо дилатированных свёрток, чтобы видеть большие области на нескольких разрешениях, и обходные связи (short-cut connections), которые проносят признаки полного разрешения мимо слоёв пониженного; без них модель не обучается.
  • Dropout на остаточном пути. Без него функция потерь на обучении падает ниже 2.0 бита на субпиксель, а на тесте поднимается выше 6.0.

PixelCNN++ достигает 2.92 бита на размерность на CIFAR-10. Та же модель, обученная на изображениях с добавленным равномерным шумом, то есть с непрерывной плотностью вместо округления, даёт оценку 3.11.

Сравнение правдоподобий#

Отрицательное логарифмическое правдоподобие на тестовой части CIFAR-10 в битах на размерность (van den Oord et al., 2016a, 2016b; Salimans et al., 2017):

МодельБит на размерность
Равномерное распределение8.00
Многомерное нормальное4.70
NICE (поток)4.48
Real NVP (поток)3.49
PixelCNN3.14
Row LSTM3.07
Gated PixelCNN3.03
Diagonal BiLSTM (PixelRNN)3.00
PixelCNN++2.92

Несколько лет после 2016 года лучшие правдоподобия на этих бенчмарках держали авторегрессионные модели. Правдоподобие измеряет, насколько хорошо модель сжимает тестовую выборку, а это не то же самое, насколько хорошо выглядят её сэмплы: это отдельно измеряют заметки об оценке, и модель может выигрывать по одному и проигрывать по другому.

Сильные стороны и ограничения#

  • Точное правдоподобие и простое обучение. Функция потерь — сумма кросс-энтропий без оценок снизу, без противника и без якобиана, и её можно отслеживать на отложенных данных.
  • Медленное сэмплирование. Одно вычисление сети на позицию: nn на изображение у модели, работающей по субпикселям, и одно на пиксель (n/3n/3 для цветного изображения), когда сеть, как PixelCNN++, предсказывает пиксели целиком, а значения каналов затем вытягиваются из её выхода. В обоих случаях стоимость растёт с разрешением. Диффузии тоже нужно много проходов, но их число задаётся расписанием шума, а не числом пикселей.
  • Нет латентного кода. Ничто не описывает изображение несколькими числами, а порядок произволен: модель должна выучить глобальную структуру из предсказания слева направо и сверху вниз. PixelVAE (Gulrajani et al., 2017) соединяет две идеи: это VAE, декодер которого — PixelCNN. Ему нужно гораздо меньше авторегрессионных слоёв, чем одному PixelCNN, а его коды компактнее, чем у обычного VAE, и всё же передают большую часть структуры изображения.
  • Сильные априорные распределения. На небольшой сетке дискретных кодов авторегрессионная модель — отличное априорное распределение. Такова её роль в VQ-VAE, где PixelCNN обучается на индексах кодов, и с этого начинается следующая заметка.

Ноутбук#

Ноутбук о PixelCNN из репозитория HSE-LAMBDA Deep Generative Models обучает небольшой PixelCNN на постерах фильмов (номера ячеек считают все ячейки, начиная с 1). Сохранённых выводов в нём нет, поэтому он описывает процедуру, а не результаты.

  • Ячейки 1–3: подготовка. Постеры размера 268×182268 \times 182 уменьшаются в 3 раза до 89×6089 \times 60; батчи по 32.
  • Ячейки 4–5, задание 1: bucketize переводит 0–255 в 10 уровней целочисленным делением на 26 с ограничением сверху числом 9. Корзины не совсем равные, как утверждает docstring: в уровнях 0–8 по 26 значений, а в уровне 9 — 22 значения 234–255. В тексте задания форма указана как batch × channels × width × height; у тензоров она batch × channels × height × width.
  • Ячейки 6–7, задание 2: изменение размера, перевод в серое, ToTensor, умножение на 255 и bucketize. Произведение здесь возвращает точные целые, так что ни один пиксель не попадает не в ту корзину.
  • Ячейки 9–13: CustomImageDataset составляет список файлов изображений в ../data/movie-genre-from-its-poster/posters/ — так называется папка датасета Kaggle Movie Genre from its Poster с постерами из IMDB. Он один раз открывает каждый файл методом verify() из PIL и пропускает те, которые не удаётся прочитать, переводит каждое изображение в RGB перед преобразованиями и возвращает с каждым изображением фиктивную метку 0, чтобы работал стандартный загрузчик. Последние 100 файлов в отсортированном порядке — валидация. Валидационный загрузчик отбрасывает последний неполный батч, так что используются 96 из них.
  • Ячейки 15–19, задание 3: make_mask и MaskedCNN, который на каждом прямом проходе умножает weight.data на маску на месте. Это работает: замаскированные веса всё равно получают градиенты, и оптимизатор сдвигает их между шагами, но следующий прямой проход снова их обнуляет — тот же эффект, что обнуление после каждого обновления в статье о PixelRNN. Умножение в прямом проходе без .data, как в разделе с кодом, оставляет их градиенты нулевыми.
  • Ячейки 20–21, задание 4: модель из одного слоя с маской A, одного с маской B и выхода 1×11 \times 1 на 10 логитов на пиксель. Два маскированных слоя видят немного: каждое предсказание дотягивается не выше чем на 2⌊k/2⌋2 \lfloor k/2 \rfloor строк вверх, 6 из 89 строк для ядер 7×77 \times 7. Модель получает индексы уровней 0–9 как float без масштабирования.
  • Ячейки 23–24: Adam, кросс-энтропия по 10 уровням и 50 эпох. После каждой эпохи два валидационных постера достраиваются со строки 10 и столбца 10: область ниже и правее обнуляется и заполняется заново пиксель за пикселем, 79⋅50=395079 \cdot 50 = 3950 проходов сети на область, а пиксели левее столбца 10 в нижних строках остаются настоящими. Функция потерь — в натах на пиксель, равномерное угадывание даёт ln⁡10≈2.30\ln 10 \approx 2.30; она измеряет задачу с 10 уровнями и несравнима с битами на размерность по 256 значениям.
  • Ячейка 25: восемь достраиваний со строки 30 и столбца 30, 59⋅30=177059 \cdot 30 = 1770 проходов сети на батч.
  • Ячейка 26: вопрос о цветных изображениях, ответ — в разделе Цветовые каналы выше.

Код#

PixelCNN для серых изображений с levels значениями на пиксель на PyTorch и проверка его рецептивного поля:

import torch
import torch.nn as nn
import torch.nn.functional as F

def make_mask(include_center, height, width):
    mask = torch.zeros(height, width)
    mask[: height // 2] = 1                       # строки выше центра
    mask[height // 2, : width // 2] = 1           # слева от центра в его строке
    if include_center:
        mask[height // 2, width // 2] = 1
    return mask

class MaskedConv2d(nn.Conv2d):
    def __init__(self, include_center, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.register_buffer('mask', make_mask(include_center, *self.kernel_size))

    def forward(self, x):
        return F.conv2d(x, self.weight * self.mask, self.bias, padding=self.padding)

class PixelCNN(nn.Module):
    def __init__(self, levels=10, hidden=64, layers=4, k=3):
        super().__init__()
        self.levels = levels
        blocks = [MaskedConv2d(False, 1, hidden, k, padding=k // 2), nn.ReLU()]
        for _ in range(layers - 1):
            blocks += [MaskedConv2d(True, hidden, hidden, k, padding=k // 2), nn.ReLU()]
        self.net = nn.Sequential(*blocks, nn.Conv2d(hidden, levels, 1))

    def forward(self, x):                         # x: [B, 1, H, W], целые уровни
        return self.net(x.float() / (self.levels - 1))   # логиты: [B, levels, H, W]

    @torch.no_grad()
    def sample(self, x, start=(0, 0)):            # заполнить x в порядке развёртки начиная со start
        x = x.clone()
        H, W = x.shape[-2:]
        for i in range(start[0], H):
            for j in range(start[1] if i == start[0] else 0, W):
                probs = self(x)[:, :, i, j].softmax(dim=1)
                x[:, 0, i, j] = torch.multinomial(probs, 1).squeeze(1)
        return x

# шаг обучения: loss = F.cross_entropy(model(x), x.long().squeeze(1)), в натах на пиксель

torch.manual_seed(0)
model = PixelCNN()
x = torch.randint(0, 10, (1, 1, 11, 11)).float().requires_grad_(True)
model(x)[0, :, 5, 5].sum().backward()             # какие входы влияют на выход в (5, 5)?
print((x.grad[0, 0, 1:6, 1:10] != 0).int())

Напечатанное окно покрывает четыре строки над позицией (5,5)(5, 5), её собственную строку и по четыре столбца в каждую сторону:

tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1],
        [1, 1, 1, 1, 1, 1, 1, 1, 0],
        [1, 1, 1, 1, 1, 1, 1, 0, 0],
        [1, 1, 1, 1, 1, 1, 0, 0, 0],
        [1, 1, 1, 1, 0, 0, 0, 0, 0]], dtype=torch.int32)

Ни у одного входа в позиции (5,5)(5, 5) или после неё в порядке развёртки нет ненулевого градиента, а нули в строках выше — это 6 пикселей слепого пятна с рисунка.

Дальше#

PixelCNN предсказывает пиксель по фиксированной окрестности. Авторегрессионные трансформеры заменяют свёртки self-attention, которое на каждом слое видит всё прошлое, и для больших изображений переходят от пикселей к дискретным кодам VQ-автоэнкодера: сетка из 256–1024 токенов описывает изображение 256×256256 \times 256, и её можно продолжать по текстовому запросу.

Литература#