stylegan2 pytorch

Simple Stylegan2 для pytorch

Простая реализация Pytorch stylegan2 на основе https://arxiv.org/abs/1912.04958, которая может быть полностью обучена из командной строки, кодирование не требуется.

Ниже приведены некоторые цветы, которых не существует.

Также эти руки

Ни эти города

Ни эти знаменитости (обученные @yoniker)

Установить

Вам понадобится машина с установленным графическим процессором и CUDA. Затем PIP установите пакет как этот

$ pip install stylegan2_pytorch

Если вы используете машину Windows, сообщается, что следующие команды работают.

$ conda install pytorch torchvision -c python
$ pip install stylegan2_pytorch

Использовать

$ stylegan2_pytorch --data /path/to/images

Вот и все. Образцы изображений будут сохранены на results/default , а модели будут периодически сохраняться в models/default .

Расширенное использование

Вы можете указать название вашего проекта с помощью

$ stylegan2_pytorch --data /path/to/images --name my-project-name

Вы также можете указать местоположение, где должны храниться промежуточные результаты и модели контрольных точек

$ stylegan2_pytorch --data /path/to/images --name my-project-name --results_dir /path/to/results/dir --models_dir /path/to/models/dir

Вы можете увеличить емкость сети (которая по умолчанию до 16 ) для улучшения результатов генерации за счет большей памяти.

$ stylegan2_pytorch --data /path/to/images --network-capacity 256

По умолчанию, если обучение будет отключено, оно автоматически возобновится из последнего контрольного файла. Если вы хотите перезапустить с новыми настройками, просто добавьте new флаг

$ stylegan2_pytorch --new --data /path/to/images --name my-project-name --image-size 512 --batch-size 1 --gradient-accumulate-every 16 --network-capacity 10

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

$ stylegan2_pytorch  --generate

Чтобы сгенерировать видео интерполяции через две случайные точки в скрытом пространстве.

$ stylegan2_pytorch --generate-interpolation --interpolation-num-steps 100

Чтобы сохранить каждую отдельную кадр интерполяции

$ stylegan2_pytorch --generate-interpolation --save-frames

Если предыдущая контрольная точка содержала лучший генератор (который часто случается, когда генераторы начинают деградировать к концу обучения), вы можете загрузить с предыдущей контрольной точки с другим флагом

$ stylegan2_pytorch --generate --load-from {checkpoint number}

Техника, используемая как в стиле, так и в Биггане, усекает скрытые значения, так что их значения приближаются к среднему. Чем мало значения усечения, тем лучше выбороны будут отображаться за счет разнообразия выборки. Вы можете контролировать это с помощью --trunc-psi , где значения обычно падают от 0.5 до 1 . Он установлен на 0.75 по умолчанию

$ stylegan2_pytorch --generate --trunc-psi 0.5

Обучение мульти-GPU

Если у вас есть одна машина с несколькими графическими процессорами, репозиторий предлагает способ использовать все их для обучения. С несколькими графическими процессорами каждая партия будет равномерно разделена между доступными графическими процессорами. Например, для 2 графических процессоров, с размером партии 32, каждый графический процессор увидит 16 образцов.

Вам просто нужно добавить флаг --multi-gpus , о том, как позаботится все остальное. Если вы хотите ограничить определенные графические процессоры, вы можете использовать переменную среды CUDA_VISIBLE_DEVICES для управления тем, какие устройства можно использовать. (Ex. CUDA_VISIBLE_DEVICES=0,2,3 только устройства 0, 2, 3 доступны)

$ stylegan2_pytorch --data ./data --multi-gpus --batch-size 32 --gradient-accumulate-every 1

Низкие объемы обучающих данных

В прошлом Гансу нужно было много данных, чтобы научиться хорошо генерировать. В качестве примера модель Faces взяла 70 тыс. Высококачественных изображений от Flickr.

Однако в мае 2020 года исследователи по всему миру независимо сходились по простой методике, чтобы уменьшить это число до 1-2K . Эта простая идея состояла в том, чтобы дифферентно увеличить все изображения, сгенерированные или реальные, входящие в дискриминатор во время обучения.

Если бы кто -то увеличивал с достаточно низкой вероятностью, увеличение не будет «протекать» в поколения.

В настройке низких данных вы можете использовать эту функцию с помощью простого флага.

 # find a suitable probability between 0. -> 0.7 at maximum
$ stylegan2_pytorch --data ./data --aug-prob 0.25

По умолчанию используемые увеличения являются translation и cutout . Если вы хотите добавить color , вы можете сделать это с аргументом --aug-types .

 # make sure there are no spaces between items!
$ stylegan2_pytorch --data ./data --aug-prob 0.25 --aug-types [translation,cutout,color]

Вы можете настроить его на любую комбинацию трех, которые вы хотели бы. Код дифференцируемого увеличения был скопирован и слегка изменен отсюда.

Когда я прекратил тренироваться?

Как можно дольше, пока не развалится состязательная игра между двумя нейронными сетями (мы называем это расхождение). По умолчанию количество обучающих этапов устанавливается на 150000 для изображений 128x128, но вы, безусловно, хотите, чтобы это число было выше, если GAN не расходится к концу тренировок, или если вы тренируетесь с более высоким разрешением.

$ stylegan2_pytorch --data ./data --image-size 512 --num-train-steps 1000000

Внимание

Эта структура также позволяет вам добавить эффективную форму самопринятого к обозначенным слоям дискриминатора (и симметричного слоя генератора), что значительно улучшит результаты. Чем больше внимания вы можете себе позволить, тем лучше!

 # add self attention after the output of layer 1
$ stylegan2_pytorch - - data . / data - - attn - layers 1 

 # add self attention after the output of layers 1 and 2
# do not put a space after the comma in the list!
$ stylegan2_pytorch - - data . / data - - attn - layers [ 1 , 2 ]

Бонус

Обучение прозрачным изображениям

$ stylegan2_pytorch --data ./transparent/images/path --transparent

Соображения памяти

Чем больше у вас памяти GPU, тем больше и лучше генерация изображений. Nvidia рекомендовала иметь до 16 ГБ для обучения изображений 1024x1024. Если у вас меньше этого, есть пара настроек, с которыми вы можете играть, чтобы модель подходила.

$ stylegan2_pytorch --data /path/to/data 
    --batch-size 3 
    --gradient-accumulate-every 5 
    --network-capacity 16

1. Размер партии-вы можете уменьшить batch-size до 1, но вы должны увеличить gradient-accumulate-every соответственно, чтобы мини-партия, которую видит сеть, не слишком мала. Это может быть запутанным для непрофессионала, поэтому я подумаю о том, как я бы автоматизировал выбор gradient-accumulate-every в будущем.

2. Емкость сети - вы можете уменьшить емкость нейронной сети, чтобы уменьшить требования к памяти. Просто имейте в виду, что это было показано, что это ухудшает производительность генерации.

Если ничего из этого не работает, вы можете согласиться на «легкий» Gan, который позволит вам компромисс, чтобы тренироваться с большими резолюциями за разумное количество времени.

Развертывание на AWS

Ниже приведены некоторые шаги, которые могут быть полезны для развертывания с использованием веб -сервисов Amazon. Чтобы использовать это, вам придется предоставить экземпляр EC2, поддерживаемый графическим процессором. Соответствующий тип экземпляра будет из серии P2 или P3. Я (Iboates) попробовал P2.xlarge (самый дешевый вариант), и это было довольно медленно, медленнее, чем использование Google Colab. Более мощные типы экземпляров могут быть лучше, но они дороже. Вы можете прочитать больше о них здесь.

Настройки шагов

1. Архивируйте свои учебные данные и загрузите их в ведро S3
2. Предоставьте ваш экземпляр EC2 (я использовал Ubuntu AMI)
3. Войдите в экземпляр EC2 через SSH
4. Установите клиент AWS CLI и настройте его:

sudo snap install aws-cli --classic
aws configure

Затем вам придется ввести ключи AWS Access, которые вы можете получить из консоли управления в консоли управления AWS> Профиль> Мои учетные данные>

Затем запустите эти команды, или, возможно, поместите их в сценарий оболочки и выполните это:

mkdir data
curl -O https://bootstrap.pypa.io/get-pip.py
sudo apt-get install python3-distutils
python3 get-pip.py
pip3 install stylegan2_pytorch
export PATH= $PATH :/home/ubuntu/.local/bin
aws s3 sync s3:// < Your bucket name > ~ /data
cd data
tar -xf ../train.tar.gz

Теперь вы должны быть в состоянии тренироваться, просто просто вызовов stylegan2_pytorch [args] .

Примечания:

• Если у вас много учебных данных, вам может потребоваться предоставить дополнительное блочное хранилище через EBS.
• Кроме того, вам может потребоваться распространять ваши данные по нескольким архивам.
• Вы должны запустить это в окне screen , чтобы оно не завершилось, как только вы выходите из сеанса SSH.

Исследовать

FID счета

Благодаря GetSeclectic, теперь вы можете периодически рассчитать оценку FID! Опять же, супер простой с одним дополнительным аргументом, как показано ниже.

Во -первых, установите пакет pytorch_fid

$ pip install pytorch-fid

С последующим

$ stylegan2_pytorch --data ./data --calculate-fid-every 5000

Результаты FID будут зарегистрированы в ./results/{name}/fid_scores.txt

Кодирование

Если вы хотите попробовать изображения программно, вы можете сделать это со следующим простым классом ModelLoader .

 import torch
from torchvision . utils import save_image
from stylegan2_pytorch import ModelLoader

loader = ModelLoader (
    base_dir = '/path/to/directory' ,   # path to where you invoked the command line tool
    name = 'default'                   # the project name, defaults to 'default'
)

noise   = torch . randn ( 1 , 512 ). cuda () # noise
styles  = loader . noise_to_styles ( noise , trunc_psi = 0.7 )  # pass through mapping network
images  = loader . styles_to_images ( styles ) # call the generator on intermediate style vectors

save_image ( images , './sample.jpg' ) # save your images, or do whatever you desire

Регистрация в экспериментирование трекера

Чтобы зарегистрировать потери на трекер эксперимента с открытым исходным кодом (AIM), вам просто нужно пропустить дополнительный флаг, как это.

$ stylegan2_pytorch --data ./data --log

Затем вам нужно убедиться, что у вас установлен Docker. Следуя инструкциям в AIM, вы выполняете следующее в своем терминале.

$ aim up

Затем откройте свой браузер по адресу, и вы должны увидеть

Экспериментальный

Top-K Training для генератора

Новая статья дала доказательства того, что, просто исказив градиентный вклад от образцов, которые считаются поддельными дискриминатором, генератор учится значительно лучше, достигая нового состояния искусства.

$ stylegan2_pytorch - - data . / data - - top - k - training

Гамма - это график распада, который медленно уменьшает топк от полного размера партии до целевой доли 50% (также модифицируемый гиперпараметт).

$ stylegan2_pytorch - - data . / data - - top - k - training - - generate - top - k - frac 0.5 - - generate - top - k - gamma 0.99

Квантование признаков

Недавняя статья сообщила о улучшенных результатах, если промежуточные представления дискриминатора являются кеккторами квантованы. Хотя я не заметил никаких драматических изменений, я решил добавить это в качестве функции, чтобы другие умы могли исследовать. Для использования вы должны указать, какие слой (ы) вы хотели бы вектор квантовать. Размер словаря по умолчанию составляет 256 и также настраиваемо.

 # feature quantize layers 1 and 2, with a dictionary size of 512 each
# do not put a space after the comma in the list!
$ stylegan2_pytorch - - data . / data - - fq - layers [ 1 , 2 ] - - fq - dict - size 512

Регуляризация контрастной потери

Я пробовал контрастное обучение на дискриминаторе (на шаг с обычным обучением GAN) и, возможно, наблюдал улучшенную стабильность и качество окончательных результатов. Вы можете включить эту экспериментальную функцию с помощью простого флага, как показано ниже.

$ stylegan2_pytorch - - data . / data - - cl - reg

Релятивистская потеря дискриминатора

Это было предложено в релятивистской газете Gan для стабилизации обучения. У меня были смешанные результаты, но я буду включать эту функцию для тех, кто хочет поэкспериментировать с ней.

$ stylegan2_pytorch - - data . / data - - rel - disc - loss

Блок 4x4

По умолчанию стиль стиля архитектуры постоянно изучил 4x4 блок, поскольку он постепенно поднимается. Это экспериментальная особенность, которая делает его таким образом, что блок 4x4 изучается из вектора стиля w вместо этого.

$ stylegan2_pytorch - - data . / data - - no - const

Двойная контрастная потеря

Недавняя статья предположила, что новая контрастная потеря между реальными и поддельными логитами может улучшить качество по сравнению с другими видами потерь. (По умолчанию в этом репозитории - потеря шарнира, и в статье наблюдается небольшое улучшение)

$ stylegan2_pytorch - - data . / data - - dual - contrast - loss 

Альтернативы

Stylegan2 + unet discinator

Я получил действительно хорошие результаты с дискриминатором UNET, но архитектурные изменения были слишком большими, чтобы соответствовать варианту в этом репозитории. Если вы стремитесь к совершенству, не стесняйтесь попробовать его.

Если вы хотите, чтобы я дал королевскую обработку какой -то другой архитектуре Gan (Biggan), не стесняйтесь обращаться по электронной почте. Рад слышать ваш шаг.

Признательность

Спасибо Мэтью Манну за его вдохновляющий простой порт для Tensorflow 2.0

Ссылки

 @article { Karras2019stylegan2 ,
    title   = { Analyzing and Improving the Image Quality of {StyleGAN} } ,
    author  = { Tero Karras and Samuli Laine and Miika Aittala and Janne Hellsten and Jaakko Lehtinen and Timo Aila } ,
    journal = { CoRR } ,
    volume  = { abs/1912.04958 } ,
    year    = { 2019 } ,
}

 @misc { zhao2020feature ,
    title   = { Feature Quantization Improves GAN Training } ,
    author  = { Yang Zhao and Chunyuan Li and Ping Yu and Jianfeng Gao and Changyou Chen } ,
    year    = { 2020 }
}

 @misc { chen2020simple ,
    title   = { A Simple Framework for Contrastive Learning of Visual Representations } ,
    author  = { Ting Chen and Simon Kornblith and Mohammad Norouzi and Geoffrey Hinton } ,
    year    = { 2020 }
}

 @article {,
    title     = { Oxford 102 Flowers } ,
    author    = { Nilsback, M-E. and Zisserman, A., 2008 } ,
    abstract  = { A 102 category dataset consisting of 102 flower categories, commonly occuring in the United Kingdom. Each class consists of 40 to 258 images. The images have large scale, pose and light variations. }
}

 @article { afifi201911k ,
    title   = { 11K Hands: gender recognition and biometric identification using a large dataset of hand images } ,
    author  = { Afifi, Mahmoud } ,
    journal = { Multimedia Tools and Applications }
}

 @misc { zhang2018selfattention ,
    title   = { Self-Attention Generative Adversarial Networks } ,
    author  = { Han Zhang and Ian Goodfellow and Dimitris Metaxas and Augustus Odena } ,
    year    = { 2018 } ,
    eprint  = { 1805.08318 } ,
    archivePrefix = { arXiv }
}

 @article { shen2019efficient ,
    author    = { Zhuoran Shen and
               Mingyuan Zhang and
               Haiyu Zhao and
               Shuai Yi and
               Hongsheng Li } ,
    title     = { Efficient Attention: Attention with Linear Complexities } ,
    journal   = { CoRR } ,  
    year      = { 2018 } ,
    url       = { http://arxiv.org/abs/1812.01243 } ,
}

 @article { zhao2020diffaugment ,
    title   = { Differentiable Augmentation for Data-Efficient GAN Training } ,
    author  = { Zhao, Shengyu and Liu, Zhijian and Lin, Ji and Zhu, Jun-Yan and Han, Song } ,
    journal = { arXiv preprint arXiv:2006.10738 } ,
    year    = { 2020 }
}

 @misc { zhao2020image ,
    title  = { Image Augmentations for GAN Training } ,
    author = { Zhengli Zhao and Zizhao Zhang and Ting Chen and Sameer Singh and Han Zhang } ,
    year   = { 2020 } ,
    eprint = { 2006.02595 } ,
    archivePrefix = { arXiv }
}

 @misc { karras2020training ,
    title   = { Training Generative Adversarial Networks with Limited Data } ,
    author  = { Tero Karras and Miika Aittala and Janne Hellsten and Samuli Laine and Jaakko Lehtinen and Timo Aila } ,
    year    = { 2020 } ,
    eprint  = { 2006.06676 } ,
    archivePrefix = { arXiv } ,
    primaryClass = { cs.CV }
}

 @misc { jolicoeurmartineau2018relativistic ,
    title   = { The relativistic discriminator: a key element missing from standard GAN } ,
    author  = { Alexia Jolicoeur-Martineau } ,
    year    = { 2018 } ,
    eprint  = { 1807.00734 } ,
    archivePrefix = { arXiv } ,
    primaryClass = { cs.LG }
}

 @misc { sinha2020topk ,
    title   = { Top-k Training of GANs: Improving GAN Performance by Throwing Away Bad Samples } ,
    author  = { Samarth Sinha and Zhengli Zhao and Anirudh Goyal and Colin Raffel and Augustus Odena } ,
    year    = { 2020 } ,
    eprint  = { 2002.06224 } ,
    archivePrefix = { arXiv } ,
    primaryClass = { stat.ML }
}

 @misc { yu2021dual ,
    title   = { Dual Contrastive Loss and Attention for GANs } ,
    author  = { Ning Yu and Guilin Liu and Aysegul Dundar and Andrew Tao and Bryan Catanzaro and Larry Davis and Mario Fritz } ,
    year    = { 2021 } ,
    eprint  = { 2103.16748 } ,
    archivePrefix = { arXiv } ,
    primaryClass = { cs.CV }
}