Träna GAN-modeller

Djupinlärning för bilder med PyTorch

Michal Oleszak

Machine Learning Engineer

Generatorns mål

GAN-arbetsflödesdiagram.

Mål: Generera falska bilder som lurar diskriminatorn
Idé: Använd diskriminatorn för att utvärdera generatorns prestanda
Generatorns utdata klassificeras av diskriminatorn som:
- Äkta (etikett 1) – bra, liten förlust
- Falsk (etikett 0) – dåligt, stor förlust

Generatorförlust

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

Definiera slumpmässigt brus
Generera en falsk bild
Hämta diskriminatorns prediktion för den falska bilden
Använd BCE-kriteriet (binär korssviktsentropy)
Generatorförlust: BCE mellan diskriminatorns prediktioner och en tensor av ettor

Diskriminatorns mål

GAN-arbetsflödesdiagram.

Mål: Korrekt klassificera falska och äkta bilder
Generatorns utdata ska klassificeras som falsk (etikett 0)
Äkta bilder ska klassificeras som äkta (etikett 1)

Diskriminatorförlust

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

Definiera BCE-kriteriet
Generera indatabrus till generatorn
Generera falska bilder
Hämta diskriminatorns prediktioner för falska bilder
Beräkna förlustkomponenten för falska bilder
Hämta diskriminatorns prediktioner för äkta bilder
Beräkna förlustkomponenten för äkta bilder
Slutlig förlust är medelvärdet av de två förlustkomponenterna

GAN-träningsloop

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()

Iterera över epoker och databatchar och beräkna aktuell batchstorlek
Nollställ diskriminatoroptimerarens gradienter
Beräkna diskriminatorförlusten
Beräkna diskriminatorns gradienter och utför optimeringssteget
Nollställ generatoroptimerarens gradienter
Beräkna generatorförlusten
Beräkna generatorns gradienter och utför optimeringssteget

Nu kör vi en övning!

Djupinlärning för bilder med PyTorch