Colorization Transformer

Исходный код, сопровождающий трансформатор раскрашивания бумаги, который будет представлен на ICLR 2021. Работа Маноджа Кумара, Дирка Вайсенборна и Нала Кальчбреннера.

Резюме статьи

Колтран состоит из трех компонентов: а colorizer, color upsamplerи spatial upsampler.

• Это colorizerавторегрессионная архитектура, основанная на самоконтроле, состоящая из условных трансформаторных слоев. Он грубо раскрашивает 64x64изображения в оттенках серого низкого разрешения пиксель за пикселем.
• Это color upsamplerпараллельная, детерминированная сеть, основанная на самоконтроле. Он уточняет грубое изображение с низким разрешением в изображение 64x64RGB.
• Архитектура spatial upsamplerтого похожа на color upsamplerсвою . Он суперразрешает изображение RGB с низким разрешением в конечный результат.
• Он colorizerимеет вспомогательную параллельную модель прогнозирования цвета, состоящую из одного линейного слоя.

Мы сообщаем результаты после обучения каждого из этих отдельных компонентов на 4x4 TPUv2чипах. Пожалуйста, отрегулируйте размер модели и размер партии во время обучения с использованием меньшего количества ресурсов. Наши результаты по обучению этих компонентов с использованием меньших ресурсов доступны в приложении.

Полные конфигурации, используемые для обучения моделей в статье, доступны в справочнике configs. Конфигурации для чрезвычайно малых моделей предоставляются в test_configs, чтобы проверить, что модель строится быстро. Установите флаг --steps_per_summaries=100 для быстрого вывода журналов. При выборке установите config.sample.log_dir в соответствующий каталог записи.

Требования

pip install -r requirements.txt 

Обучение

Выполните следующую команду, чтобы обучить колоризатор

python -m coltran.run --config=coltran/configs/colorizer.py --mode=train --logdir=/colorizer_ckpt_dir 

Чтобы обучить цветовой и пространственный апсамплер, замените configs/colorizer.py на configs/color_upsampler.py и configs/spatial_upsampler.py соответственно

Оценка

Для оценки,

python -m coltran.run --config=coltran/configs/colorizer.py --mode=eval_valid --logdir=/colorizer_ckpt_dir 

Отбор проб

Одиночная выборка GPU

Выборка изображений с высоким разрешением-это трехэтапная процедура. На графическом процессоре P100 колоризатор сэмплирует партию из 20 изображений за 3-5 минут, в то время как цвет и пространственный апсамплер сэмплируют в порядке миллисекунд.

Конфигурации выборки для каждой модели описываются config.sampleConfigDict at configs/.py

• sample.num_outputs - общее количество изображений в оттенках серого
• sample.logdir - пример каталога записи.
• sample.gen_data_dir - путь к тому месту, где хранятся образцы с предыдущего шага.
• sample.skip_batches - пропускаются первые skip_batches*batch_sizeизображения из общедоступных TF-наборов данных imagenet.

Пожалуйста, убедитесь, что num_outputsи skip_batchesявляются одинаковыми во всех трех компонентах. Сгенерированные образцы записываются в виде TFRecords для $logdir/${config.sample.logdir}

Краситель

Команда сэмплирует грубые цветные изображения 64x64 с низким разрешением.

python -m coltran.sample --config=coltran/configs/colorizer.py --mode=sample_test --logdir=/colorizer_ckpt_dir 

Цветной Апсамплер

Команда преобразует грубые изображения 64x64 из предыдущего шага в более тонкие изображения 64x64.

Внимание: пожалуйста, установите config.sample.gen_data_dirконфигурацию color upsampler на /colorizer_ckpt_dir/${config.sample.logdir}

python -m coltran.sample --config=coltran/configs/color_upsampler.py --mode=sample_test --logdir=/cup_ckpt_dir 

Пространственный Апсемплер

Команда superresolves предыдущий вывод на выход с высоким разрешением 256x256.

Внимание: пожалуйста, установите config.sample.gen_data_dirконфигурацию пространственного upsampler на $/cup_ckpt_dir/${config.sample.logdir}

python -m coltran.sample --config=coltran/configs/spatial_upsampler.py --mode=sample_test --logdir=/cup_ckpt_dir 

Выборка с несколькими графическими процессорами

Выборка может быть распараллелена между пакетами в установке с несколькими графическими процессорами с помощью флага config.sample.skip_batches. Например, в установке с 2 машинами и размером партии 20, чтобы раскрасить 100 оттенков серого изображения на машинный набор config.sample.skip_batchesпервой и второй машины до 0 и 5 соответственно.

Предварительно подготовленные контрольно-пропускные пункты

Мы выпускаем предварительно подготовленные контрольные точки в ImageNet по следующему URL-адресу

• Краситель - Ссылка
• Цветной Upsampler - Link
• Пространственный Upsampler-Link

Чтобы выполнить выборку, загрузите их в локальный каталог и запустите сценарий выборки с logdirфлагом, установленным в локальном пути.

Эталонные Тензорные Доски

Сводные данные о наших поездах находятся в открытом доступе через tensorboard.dev

• Краситель - Ссылка
• Цветной Upsampler - Link
• Пространственный Upsampler-Link

Разбор TFRecords

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

def parse_example(example_proto, res=64):   features = {'image': tf.io.FixedLenFeature([res*res*3], tf.int64)}   example = tf.io.parse_example(example_proto, features=features)   image = tf.reshape(example['image'], (res, res, 3))   return image  gen_dataset = tf.data.TFRecordDataset(listdir(path)) gen_dataset = gen_dataset.map(lambda x: parse_example(x, res)) gen_dataset = iter(gen_dataset) for image in gen_dataset:   plt.imshow(image) 

Цитирование

Если вы используете код или модель, пожалуйста, процитируйте нашу статью.

@inproceedings{ kumar2021colorization, title={Colorization Transformer}, author={Manoj Kumar and Dirk Weissenborn and Nal Kalchbrenner}, booktitle={International Conference on Learning Representations}, year={2021}, url={https://openreview.net/forum?id=5NA1PinlGFu} }  

Источник: github.com

Комментарии: