Trénování GAN

Deep Learning pro obrázky s PyTorchem

Michal Oleszak

Machine Learning Engineer

Cíl generátoru

Diagram pracovního postupu GAN.

Cíl: Generovat falešné obrázky, které oklamou diskriminátor
Myšlenka: Využít diskriminátor k hodnocení výkonu generátoru
Výstup generátoru klasifikovaný diskriminátorem jako:
- Reálný (štítek 1) – dobré, malá ztráta
- Falešný (štítek 0) – špatné, velká ztráta

Ztráta generátoru

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

Definujte náhodný šum
Generujte falešný obrázek
Získejte predikci diskriminátoru pro falešný obrázek
Použijte kritérium binární křížové entropie (BCE)
Ztráta generátoru: BCE mezi predikcemi diskriminátoru a tenzorem jedniček

Cíl diskriminátoru

Diagram pracovního postupu GAN.

Cíl: Správně klasifikovat falešné a reálné obrázky
Výstupy generátoru by měly být klasifikovány jako falešné (štítek 0)
Reálné obrázky by měly být klasifikovány jako reálné (štítek 1)

Ztráta diskriminátoru

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)

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

    disc_loss = (real_loss + fake_loss) / 2
    return disc_loss

Definujte kritérium binární křížové entropie
Generujte vstupní šum pro generátor
Generujte falešné obrázky
Získejte predikce diskriminátoru pro falešné obrázky
Vypočítejte složku ztráty pro falešné obrázky
Získejte predikce diskriminátoru pro reálné obrázky
Vypočítejte složku ztráty pro reálné obrázky
Celková ztráta je průměr složek reálné a falešné ztráty

Trénovací smyčka 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()

Iterujte přes epochy a dávky reálných dat a vypočítejte velikost aktuální dávky
Vynulujte gradienty optimalizátoru diskriminátoru
Vypočítejte ztrátu diskriminátoru
Vypočítejte gradienty diskriminátoru a proveďte optimalizační krok
Vynulujte gradienty optimalizátoru generátoru
Vypočítejte ztrátu generátoru
Vypočítejte gradienty generátoru a proveďte optimalizační krok

Pojďme si procvičit!

Deep Learning pro obrázky s PyTorchem