WGAN loss diverges

user_11239518 0
Viewed 299

I have some trouble to understand the WGAN Loss values.

I understand that we do not have a discriminator anymore, but a critic. Difference is, that the Discriminator tries to classify the input ergo map it to either 0 or 1 and that the critic tries to score them. The scoring map to any real number and the Loss function is:

L_critic = Critic(fake_input) - Critic(real_input) + weight * gradient_penalty

And the Loss function for the generator which produces the fake_input is:

L_generator = - Critic(fake_input)

My question is: Over the Epochs my Loss values for the generator and Discriminator diverges and takes really high numbers. Is this common? I have seen elsewhere that the Loss of the Critic drop to 0, but what ensures that? As far as I understood the Critic tries to score the real from fake as far apart as possible.

[Loss evaluation for the critic and generator. Y-Axis is a factor of 10^11]

loss pytorch python
2 Answers

There is no information for model has which layers and train how, in question. So I would like to share my experience on solving this problem.

I got exactly same problem and finally solved this. I tried to train with DCGAN based model (yes, it has convolutional 2d layer)

I still can not understand perfectly, but this works.

  1. Remove batch normalization for Discriminator(Critic)

  2. Train with model based MLP. NOT Convolution2D

I don't understand. when using Conv2d on any of Generator and Discriminator, loss diverges!

here are two model which works successfully and get loss fine.

# channels = 3
# image_size = 64

class Generator(nn.Module):
    def __init__(self):
        super(Generator, self).__init__()
        
        self.model = nn.Sequential(
            nn.Linear(nz, 128),
            nn.BatchNorm1d(128, 0.8),
            nn.LeakyReLU(0.2, inplace=True),
            
            nn.Linear(128, 256),
            nn.BatchNorm1d(256, 0.8),
            nn.LeakyReLU(0.2, inplace=True),
            
            nn.Linear(256, 512),
            nn.BatchNorm1d(512, 0.8),
            nn.LeakyReLU(0.2, inplace=True),
            
            nn.Linear(512, 1024),
            nn.BatchNorm1d(1024, 0.8),
            nn.LeakyReLU(0.2, inplace=True),
            
            nn.Linear(1024, channels * image_size * image_size),
            nn.Tanh()
        )

    def forward(self, z):
        img = self.model(z)
        img = img.view(img.shape[0], channels, image_size, image_size)
        return img

class Discriminator(nn.Module):
    def __init__(self):
        super(Discriminator, self).__init__()

        self.model = nn.Sequential(
            nn.Linear(channels * image_size * image_size, 512),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Linear(512, 256),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Linear(256, 1),
        )

    def forward(self, img):
        img = img.view(img.shape[0], -1) # flatten
        return self.model(img)

I hope it helps you either.

user_136775540

I came across a similar issue when I was working with WGAN the last time, to solve this issue I modified my loss a bit to train both the generator and discriminator network. You can try it as well, and it may help you in solving your issue:

Gradient penalty function

def gradient_penalty(f, real, fake=None):
    def interpolate(a, b=None):
        if b is None:  # interpolation in DRAGAN
            beta = torch.rand_like(a)
            b = a + 0.5 * a.var().sqrt() * beta
        alpha = torch.rand(a.size(0), 1, 1, 1)
        alpha = alpha.cuda() if self.gpu else alpha
        inter = a + alpha * (b - a)
        return inter
    x = interpolate(real, fake).requires_grad_(True)
    pred = f(x)
    if isinstance(pred, tuple):
        pred = pred[0]
    grad = autograd.grad(outputs=pred, inputs=x,
                grad_outputs=torch.ones_like(pred),
                create_graph=True, retain_graph=True, only_inputs=True)[0]
    grad = grad.view(grad.size(0), -1)
    norm = grad.norm(2, dim=1)
    gp = ((norm - 1.0) ** 2).mean()
    return gp

Trainer function

#training generator
img_fake = Generator(input_random_vector)
d_fake = Discriminator(img_fake)

generartor_fake_loss = -d_fake.mean()
       
generator_loss = generartor_fake_loss 
        
optimizer_Generator.zero_grad()
generator_loss.backward()
optimizer_Generator.step()

#training discriminator
d_real = Discriminator(img_real)
d_fake = Discriminator(img_fake)
wd = d_real.mean() - d_fake.mean()
df_loss = -wd
df_gp = gradient_penalty(Discriminator, img_real, img_fake)
discriminator_loss = df_loss + ALPHA * df_gp 
        
optimizer_Discriminator.zero_grad()
discriminator_loss.backward()
optimizer_Discriminator.step()

You might have to experiment with the value of ALPHA which is used to calculate discriminator_loss. This should work with diverging losses.

user_132014640
Related