Навчання GAN

Глибоке навчання для зображень із PyTorch

Michal Oleszak

Machine Learning Engineer

Мета генератора

Схема роботи GAN.

Мета: генерувати фейки, що обманюють дискримінатор
Ідея: використовувати дискримінатор, щоб оцінити роботу генератора
Вихід генератора класифікується дискримінатором як:
- Реальне (мітка 1) — добре, мала втрата
- Фейк (мітка 0) — погано, велика втрата

Втрата генератора

def gen_loss(gen, disc, num_images, z_dim):

    noise = torch.randn(num_images, z_dim)

    fake = gen(noise)

    disc_pred = disc(fake)

    criterion = nn.BCEWithLogitsLoss()

    gen_loss = criterion(
        disc_pred, torch.ones_like(disc_pred)
    )
    return gen_loss

Визначити випадковий шум
Згенерувати фейкове зображення
Отримати прогноз дискримінатора для фейка
Використати критерій бінарної крос-ентропії (BCE)
Втрата генератора: BCE між прогнозами дискримінатора та тензором одиниць

Мета дискримінатора

Схема роботи GAN.

Мета: коректно класифікувати фейки та реальні зображення
Виходи генератора мають класифікуватися як фейки (мітка 0)
Реальні зображення мають класифікуватися як реальні (мітка 1)

Втрата дискримінатора

def disc_loss(gen, disc, real, num_images, z_dim):

    criterion = nn.BCEWithLogitsLoss()

    noise = torch.randn(num_images, z_dim)

    fake = gen(noise)

    disc_pred_fake = disc(fake)

    fake_loss = criterion(
        disc_pred_fake,
        torch.zeros_like(disc_pred_fake)
    )

    disc_pred_real = disc(real)

DNT_CURLLY_TAG_7

    real_loss = criterion(
        disc_pred_real,
        torch.ones_like(disc_pred_real)
    )

    disc_loss = (real_loss + fake_loss) / 2
    return disc_loss

Визначити критерій бінарної крос-ентропії
Згенерувати вхідний шум для генератора
Згенерувати фейки
Отримати прогнози дискримінатора для фейків
Обчислити компонент втрати для фейків
Отримати прогнози дискримінатора для реальних зображень
Обчислити компонент втрати для реальних
Підсумкова втрата — середнє між компонентами для реальних і фейків

Цикл навчання GAN

for epoch in range(num_epochs):
    for real in dataloader:
        cur_batch_size = len(real)


        disc_opt.zero_grad()

        disc_loss = disc_loss(
            gen, disc, real, cur_batch_size, z_dim=16)

        disc_loss.backward()
        disc_opt.step()


        gen_opt.zero_grad()

        gen_loss = gen_loss(
            gen, disc, cur_batch_size, z_dim=16)

        gen_loss.backward() 
        gen_opt.step()

Переберіть епохи та пакети реальних даних і обчисліть розмір поточного пакета
Скиньте градієнти оптимізатора дискримінатора
Обчисліть втрату дискримінатора
Обчисліть градієнти дискримінатора та виконайте крок оптимізації
Скиньте градієнти оптимізатора генератора
Обчисліть втрату генератора
Обчисліть градієнти генератора та виконайте крок оптимізації

Давайте потренуємось!

Глибоке навчання для зображень із PyTorch