Изграждане на най-простите GAN в PyTorch

Прекарах дълго време в правене GANs в TensorFlow/Keras. Твърде дълго, честно, защото промяната е трудна. Отне малко убеждаване, но в крайна сметка прехвърлих куршума и преминах към PyTorch. За съжаление, повечето от уроците за PyTorch GAN, на които се натъкнах, бяха прекалено сложни, фокусирани повече върху теорията на GAN, отколкото върху приложението, или странно непитонични. За да поправя това, написах този микро урок за създаване на ванилов GAN в PyTorch, с акцент върху PyTorch. Самият код е достъпен тук (имайте предвид, че кодът на github и същността в този урок се различават леко). Препоръчвам да отворите този урок в два прозореца, единият с разглеждания код, а другият с обясненията.

Изисквания

  1. Python 3.7 или по-нова версия. Всяко по-ниско и ще трябва да преработите f-струните.
  2. PyTorch 1.5 Не сте сигурни как да го инсталирате? „Това може да помогне“.
  3. Седемнадесет или осемнадесет минути от вашето време. Най-малко дванадесет, ако си умен.

Предстоящата задача

Създайте функция G: Z → X, където Z~U(0, 1) и X~N(0, 1).

На английски това е „направете GAN, който приближава „нормалното“ разпределение, даден „равномерен“ случаен шум като вход“. Това означава, че входът към GAN ще бъде едно число, както и изходът. Имайте предвид, че тук ще използваме функция за генериране на данни вместо набор от данни за обучение в името на простотата.

Нека просто скочим в него

Уверете се, че имате инсталирана правилната версия на Python и инсталирайте PyTorch. След това направете нов файл vanilla_GAN.py и добавете следните импортирания:

import torch
from torch import nn
import torch.optim as optim

Нашият GAN скрипт ще има три компонента: мрежа Generator, мрежа Discriminator и самият GAN, който съдържа и обучава двете мрежи. Да започнем с генератора:

Генератор

Добавете следното към вашия скрипт:

Нашият клас Generator наследява от класа nn.Module на PyTorch, който е базовият клас за модули за невронни мрежи. Накратко, той казва на PyTorch „това е невронна мрежа“. Нашият клас генератор има два метода:

Генератор.__init__

Инициализирайте обекта. Първо, това извиква метода nn.Module __init__ с помощта на super. След това създава подмодулите (т.е. слоеве) и ги присвоява като променливи на екземпляр. Те включват:

  • Линеен (т.е. напълно свързан, т.е. плътен) слой с входна ширина latent_dim и изходна ширина 64.
  • Линеен слой с входна ширина 64 и изходна ширина 32.
  • Линеен слой с входна ширина 32 и изходна ширина 1.
  • Активирането на LeakyReLU.
  • Посоченото активиране на изхода.

Тъй като тези модули се записват като променливи на екземпляр в клас, който наследява от nn.Module, PyTorch е в състояние да ги следи, когато дойде време за обучение на мрежата; повече за това по-късно.

Генератор.напред

Методът forward е от съществено значение за всеки клас, наследяващ от nn.Module, тъй като дефинира структурата на мрежата. PyTorch използва рамка define-by-run, което означава, че изчислителната графика на невронната мрежа се изгражда автоматично, докато свързвате прости изчисления заедно. Всичко е много Pythonic. В нашия forward метод преминаваме през модулите на генератора и ги прилагаме към изхода на предишния модул, връщайки крайния изход. Когато стартирате мрежата (напр.: prediction = network(data), методът forward се извиква за изчисляване на изхода.

Генератор.назад

не! PyTorch използва Autograd за автоматично разграничаване; когато стартирате метода forward, PyTorch автоматично следи изчислителната графика и следователно не е нужно да му казвате как да разпространява обратно градиентите. Как изглежда това на практика? Продължавай да четеш.

Дискриминатор

Добавете следното към вашия скрипт:

Нашият обект Discriminator ще бъде почти идентичен с нашия генератор, но като погледнете класа, може да забележите две разлики. Първо, мрежата е параметризирана и леко преработена, за да стане по-гъвкава. Второ, изходната функция е фиксирана на Sigmoid, тъй като дискриминаторът ще има за задача да класифицира проби като реални (1) или генерирани (0).

Дискриминатор.__init__

Методът Discriminator __init__ прави три неща. Отново извиква метода nn.Module __init__, използвайки super. След това записва входното измерение като обектна променлива. Накрая извиква метода _init_layers. Като аргументи __init__ приема входно измерение и списък от цели числа, наречен layers, който описва ширините на модулите nn.Linear, включително изходния слой.

Дискриминатор._init_layers

Този метод инстанцира мрежовите модули. Тялото на този метод можеше да бъде поставено в __init__, но смятам, че е по-чисто шаблонът за инициализация на обекта да бъде отделен от кода за изграждане на модул, особено когато сложността на мрежата расте. Този метод итерира аргумента layers и създава списък с модули nn.Linear с подходящ размер, както и активации на Leaky ReLU след всеки вътрешен слой и активиране на Sigmoid след последния слой. Тези модули се съхраняват в ModuleList обект, който функционира като обикновен списък на Python, с изключение на факта, че PyTorch го разпознава като списък с модули, когато дойде време за обучение на мрежата. Има и клас ModuleDict, който служи за същата цел, но функционира като речник на Python; повече за тях по-късно.

Дискриминатор.напред

Методът forward функционира по същия начин като неговия партньор в генератора. Въпреки това, тъй като запазихме нашите модули като списък, можем просто да преминем през този списък, прилагайки всеки модул на свой ред.

VanillaGAN

Добавете следното към вашия скрипт:

Нашият клас VanillaGAN съдържа обектите Generator и Discriminator и обработва тяхното обучение.

VanillaGAN.__init__

Като вход конструкторът VanillaGAN приема:

  • Обект Генератор.
  • Обект на дискриминатор.
  • Функция за шум. Това е функцията, използвана за вземане на проби от латентни вектори Z, които нашият генератор ще картографира към генерирани проби X. Тази функция трябва да приеме цяло число num като вход и да върне 2D Torch тензор с форма (num, latent_dim).
  • Функция за данни. Това е функцията, която нашият Генератор е натоварен да научи. Тази функция трябва да приеме цяло число num като вход и да върне 2D Torch тензор с форма (num, data_dim), където data_dim е измерението на данните, които се опитваме да генерираме, input_dim на нашия дискриминатор.
  • По желание, размер на минипартида за обучение.
  • По желание, устройство. Това може да бъде cpu или cuda, ако искате да използвате GPU.
  • По желание, скорости на обучение за генератора и дискриминатора.

Където е подходящо, тези аргументи се записват като променливи на екземпляр.

Целта на GAN е „двоичната кръстосана загуба на ентропия“ (nn.BCELoss), която създаваме и присвояваме като обектна променлива criterion.

Нашият GAN използва два оптимизатора, един за генератора и един за дискриминатора. Нека разбием оптимизатора на Generator, екземпляр Adam. Оптимизаторите управляват актуализации на параметрите на невронна мрежа, като се имат предвид градиентите. За да направи това, оптимизаторът трябва да знае за кои параметри трябва да се занимава; в този случай това е discriminator.parameters(). Преди няколко минути ви казах

PyTorch е в състояние да следи [модули], когато дойде време за обучение на мрежата.

Тъй като обектът Discriminator наследява от nn.Module, той наследява метода parameters, който връща всички обучаеми параметри във всички модули, зададени като променливи на екземпляра за Discriminator (ето защо трябваше да използваме nn.ModuleList вместо списък на Python, така че PyTorch да знае проверете всеки елемент за параметри). Оптимизаторът също получава определена скорост на обучение и бета параметри, които работят добре за GAN. Оптимизаторът на генератора работи по същия начин, но вместо това следи параметрите на генератора и използва малко по-малка скорост на обучение.

И накрая, ние съхраняваме колонен вектор от единици и колонен вектор от нули като етикети на клас за обучение, така че да не се налага многократно да ги възстановяваме.

VanillaGAN.generate_samples

Това е помощна функция за получаване на произволни проби от генератора. Извикан без никакви аргументи, той генерира batch_size проби. Това може да бъде заменено чрез указване на аргумента num за създаване на num проби или чрез предоставяне на 2D PyTorch тензор, съдържащ определени латентни вектори. Контекстният мениджър no_grad казва на PyTorch да не си прави труда да следи градиентите тук, намалявайки обема на изчисленията.

VanillaGAN.train_step_generator

Тази функция изпълнява една тренировъчна стъпка на генератора и връща загубата като плаваща стойност. Заедно със стъпката за обучение на дискриминатора, това е същността на алгоритъма, така че нека преминем през него ред по ред:

self.generator.zero_grad()

Изчистете градиентите. Най-якото нещо при PyTorch е, че градиентът автоматично се натрупва във всеки параметър, докато се използва мрежата. Въпреки това, ние обикновено искаме да изчистим тези градиенти между всяка стъпка на оптимизатора; методът zero_grad прави точно това.

latent_vec = self.noise_fn(self.batch_size)

Примерни batch_size латентни вектори от функцията за генериране на шум. лесно.

generated = self.generator(latent_vec)

Подайте латентните вектори в генератора и получете генерираните проби като изход (под капака тук се извиква методът generator.forward). Не забравяйте, че PyTorch е define-by-run, така че това е точката, в която се изгражда изчислителната графика на генератора.

classifications = self.discriminator(generated)

Захранете генерираните проби в Discriminator и получете неговата увереност, че всяка проба е истинска. Не забравяйте, че Дискриминаторът се опитва да класифицира тези проби като фалшиви (0), докато Генераторът се опитва да го подмами да мисли, че са истински (1). Точно както в предишния ред, това е мястото, където се изгражда изчислителната графика на Дискриминатора и тъй като са му дадени генерираните проби generated като вход, тази изчислителна графика е заседнала в края на изчислителната графика на Генератора.

loss = self.criterion(classifications, self.target_ones)

Изчислете загубата за генератора. Нашата функция за загуба е двоична кръстосана ентропия, така че загубата за всяка от batch_size проби се изчислява и осреднява в една стойност. loss е тензор на PyTorch с една единствена стойност в него, така че все още е свързан с пълната изчислителна графика.

loss.backward()

Тук се случва магията. Или по-скоро тук се получава престижът, тъй като през цялото това време магията се случва невидимо. Тук методът backward изчислява градиента d_loss/d_x за всеки параметър x в изчислителната графика.

self.optim_g.step()

Приложете една стъпка от оптимизатора, като избутвате всеки параметър надолу по градиента. Ако сте създавали GAN в Keras преди, вероятно сте запознати с необходимостта да зададете my_network.trainable = False. Едно от предимствата на PyTorch е, че не е нужно да се занимавате с това, защото на optim_g беше казано да се занимава само с параметрите на нашия генератор.

return loss.item()

Върнете загубата. Ще ги съхраняваме в списък за по-късна визуализация. Жизненоважно е обаче да използваме метода item, за да го върнем като float, а не като тензор на PyTorch. Това е така, защото, ако запазим препратка към този тензорен обект в списък, Python също ще се придържа към цялата изчислителна графика. Това е голяма загуба на памет, така че трябва да сме сигурни, че запазваме само това, от което се нуждаем (стойността), така че събирачът на отпадъци на Python да може да изчисти останалото.

VanillaGAN.train_step_discriminator

Подобно на предишния метод, train_step_discriminator изпълнява една стъпка на обучение за дискриминатора. Нека да го разгледаме ред по няколко реда:

self.discriminator.zero_grad()

Познавате този!

# real samples
real_samples = self.data_fn(self.batch_size)
pred_real = self.discriminator(real_samples)
loss_real = self.criterion(pred_real, self.target_ones)

Извадете някои реални проби от целевата функция, вземете увереността на Дискриминатора, че те са реални (Дискриминаторът иска да максимизира това!) и изчислете загубата. Това е много подобно на стъпката на обучение на генератора.

# generated samples
latent_vec = self.noise_fn(self.batch_size)
with torch.no_grad():
    fake_samples = self.generator(latent_vec)
pred_fake = self.discriminator(fake_samples)
loss_fake = self.criterion(pred_fake, self.target_zeros)

Вземете проби от няколко генерирани проби от генератора, получете увереността на Дискриминатора, че те са реални (Дискриминаторът иска да минимизира това!) и изчислете загубата. Тъй като ние обучаваме Discriminator тук, ние не се интересуваме от градиентите в Generator и като такъв използваме no_grad контекстния мениджър. Като алтернатива можете да се откажете от no_grad и да замените в реда pred_fake = self.discriminator(fake_samples.detach()) и да отделите fake_samples от изчислителната графика на Генератора след факта, но защо да си правите труда да го изчислявате на първо място?

# combine
loss = (loss_real + loss_fake) / 2

Осреднете изчислителните графики за реалните проби и генерираните проби. Да, наистина е така. Това е любимият ми ред в целия скрипт, защото PyTorch може да комбинира и двете фази на изчислителната графика, използвайки проста аритметика на Python.

loss.backward()
self.optim_d.step()
return loss_real.item(), loss_fake.item()

Изчислете градиентите, приложете една стъпка на градиентно спускане и върнете загубите.

VanillaGAN.train_step

Този метод просто прилага една стъпка на обучение на дискриминатора и една стъпка на генератора, връщайки загубите като кортеж.

Сглобяване на всичко

Добавете следното към вашия скрипт:

Функцията main е доста разбираема, но нека я прегледаме заедно за пълнота.

  • Ние внасяме time, защото обикновено е добра идея да зададете време за обучение на невронни мрежи.
  • Ще имаме 600 епохи с 10 партиди във всяка; партидите и епохите не са необходими тук, тъй като използваме истинската функция вместо набор от данни, но нека се придържаме към конвенцията за умствено удобство.
  • Ние създаваме генератора и дискриминатора. Не забравяйте, че трябва да посочим ширината на слоя на дискриминатора.
  • Ние дефинираме функцията на шума като произволни, еднакви стойности в [0, 1], изразени като колонен вектор. Посочваме устройството като „cpu“, но това може да е „CUDA“, ако сте го настроили. Имайте предвид, че ако използвате cuda тук, използвайте го за целевата функция и VanillaGAN.
  • Ние дефинираме целевата функция като произволни, нормални (0, 1) стойности, изразени като колонен вектор. Отново посочваме устройството като „cpu“.
  • Ние създаваме VanillaGAN.
  • Ние създаваме списъците, за да следим загубите и да изпълняваме тренировъчния цикъл, като отпечатваме тренировъчни статистики след всяка епоха.

Това е! Поздравления, написахте първия си GAN в PyTorch. Не включих кода за визуализация, но ето как изглежда наученото разпределение G след всяка стъпка на обучение:

И ето загубата за епохи:

Заключителни мисли

Тъй като този урок беше за изграждането на GAN класовете и обучителния цикъл в PyTorch, малко внимание беше отделено на действителната мрежова архитектура. Съвременните „хакове на GAN“ не бяха използвани и като такова крайното разпределение само до голяма степен прилича на истинското стандартно нормално разпределение. Ако се интересувате от „да научите повече за GAN“, опитайте да промените хиперпараметрите и модулите; отговарят ли резултатите на това, което очаквате?

Обикновено нямаме достъп до истинската дистрибуция, генерираща данни (ако имахме, нямаше да имаме нужда от GAN!). В „последващ урок“ към този, ние ще внедрим конволюционен GAN, който използва реален целеви набор от данни вместо функция.

Всички изображения, които не са цитирани, са мои. Чувствайте се свободни да ги използвате, но моля, цитирайте тази статия ❤️



ConorLazarou/pytorch-generative-models
VanillaGAN — Най-простият GAN, внедрен в PyTorchgithub.com