Trenowanie GAN-ów

Głębokie uczenie dla obrazów z PyTorch

Michal Oleszak

Machine Learning Engineer

Cel generatora

Diagram przepływu pracy GAN.

Cel: generowanie fałszywych obrazów, które oszukają dyskryminator
Idea: wykorzystanie dyskryminatora do oceny działania generatora
Dane wyjściowe generatora klasyfikowane przez dyskryminator jako:
- Rzeczywiste (etykieta 1) – dobrze, mała strata
- Fałszywe (etykieta 0) – źle, duża strata

Strata generatora

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

Zdefiniuj losowy szum
Wygeneruj fałszywy obraz
Pobierz predykcję dyskryminatora dla fałszywego obrazu
Użyj kryterium binarnej entropii krzyżowej (BCE)
Strata generatora: BCE między predykcjami dyskryminatora a tensorem jedynek

Cel dyskryminatora

Diagram przepływu pracy GAN.

Cel: poprawna klasyfikacja obrazów fałszywych i rzeczywistych
Dane wyjściowe generatora powinny być klasyfikowane jako fałszywe (etykieta 0)
Obrazy rzeczywiste powinny być klasyfikowane jako rzeczywiste (etykieta 1)

Strata dyskryminatora

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

Zdefiniuj kryterium binarnej entropii krzyżowej
Wygeneruj szum wejściowy dla generatora
Wygeneruj obrazy fałszywe
Pobierz predykcje dyskryminatora dla obrazów fałszywych
Oblicz składnik straty dla obrazów fałszywych
Pobierz predykcje dyskryminatora dla obrazów rzeczywistych
Oblicz składnik straty dla obrazów rzeczywistych
Strata końcowa to średnia składników straty dla obrazów rzeczywistych i fałszywych

Pętla trenowania 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()

Iteruj po epokach i wsadach danych rzeczywistych; oblicz rozmiar bieżącego wsadu
Wyzeruj gradienty optymalizatora dyskryminatora
Oblicz stratę dyskryminatora
Oblicz gradienty dyskryminatora i wykonaj krok optymalizacji
Wyzeruj gradienty optymalizatora generatora
Oblicz stratę generatora
Oblicz gradienty generatora i wykonaj krok optymalizacji

Czas na ćwiczenia!

Głębokie uczenie dla obrazów z PyTorch