Problem Training a Wasserstein GAn with Gradient Penalty

username: Federico_Ottomano
cooked: <p>I’m looking to re-implement in <code>Pytorch</code> the following WGAN-GP model:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/1/9/19c1d85478418c3ec6758b6a17bb1a145b639916.png" data-download-href="https://discuss.pytorch.org/uploads/default/19c1d85478418c3ec6758b6a17bb1a145b639916" title="model"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/1/9/19c1d85478418c3ec6758b6a17bb1a145b639916_2_487x500.png" alt="model" data-base62-sha1="3FRfmHC3XYiSxF5TGIErobxR0Y6" width="487" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/1/9/19c1d85478418c3ec6758b6a17bb1a145b639916_2_487x500.png, https://discuss.pytorch.org/uploads/default/original/3X/1/9/19c1d85478418c3ec6758b6a17bb1a145b639916.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/1/9/19c1d85478418c3ec6758b6a17bb1a145b639916.png 2x" data-dominant-color="EBEBEB"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">model</span><span class="informations">664×681 90.1 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>taken by this <a href="https://static-content.springer.com/esm/art%3A10.1038%2Fs41524-020-00352-0/MediaObjects/41524_2020_352_MOESM1_ESM.pdf" rel="noopener nofollow ugc">paper</a>.</p>
<p>The original implementation was in <code>tensorflow</code>. Apart from minor issues which require me to modify subtle details, since torch seems not supporting <code>padding='same'</code> for strided convolutions, my implementation is the following:</p>
<pre><code class="lang-auto">class Discriminator(nn.Module):
    
    def __init__(self):
        
        super(Discriminator, self).__init__()
        
        self.disc = nn.Sequential(
            
            nn.Conv2d(in_channels=1, out_channels=32, kernel_size = 3, stride = (1, 1),padding='same'),
            self._block(in_channels=32, out_channels=32, kernel_size=3, stride=(2,1), padding=(1,1)),
            self._block(in_channels=32, out_channels=64, kernel_size = 3, stride = (1, 1),padding='same'),
            self._block(in_channels=64, out_channels=64, kernel_size = 3, stride = (2, 1),padding=(1,1)),
            self._block(in_channels=64, out_channels=128, kernel_size = 3, stride = (1, 1),padding='same'),
            self._block(in_channels=128, out_channels=128, kernel_size = 3, stride = (2, 1),padding=(1,1)),
            self._block(in_channels=128, out_channels=256, kernel_size=5, stride=(2,2),padding=(2,2))
            )
        
        self.lin = nn.Linear(256*6*4,1)
        
    #unifies Conv2d leakyrelu and batchnorm
    def _block(self, in_channels,
              out_channels, 
              kernel_size, stride, padding):
        
        return nn.Sequential(nn.Conv2d(in_channels,
                                       out_channels,
                                       kernel_size,
                                       stride,
                                       padding,
                                       bias=False),
                nn.BatchNorm2d(out_channels),
                nn.LeakyReLU(0.2)) #bias false as we use batchnorm
    
    def forward(self, x):
        
        x = self.disc(x)
        
        x = x.view(-1,256*6*4)
        
        
        return self.lin(x)

    
class Generator(nn.Module):
    
    def __init__(self, z_dim):
        
        super(Generator, self).__init__()
        
        self.z_dim = z_dim
        
        self.lin1 = nn.Linear(z_dim, 6*4*256)
        
        self.gen = nn.Sequential(
            
            self._block(in_channels=256, out_channels=128, kernel_size=(5,4),stride=(2,2),padding=(2,1)),
            self._block(in_channels=128, out_channels=128, kernel_size=(4,3), stride=(2,1),padding=(1,1)),
            self._block(in_channels=128, out_channels=64, kernel_size=(3,3), stride=(1,1), padding=(1,1)),
            self._block(in_channels=64, out_channels=64, kernel_size=(3,3), stride=(1,1), padding=(1,1)),
            self._block(in_channels=64, out_channels=64, kernel_size=(3,2), stride=(2,2), padding=(1,4)),
            self._block(in_channels=64, out_channels=32, kernel_size=(3,3), stride=(1,1), padding=(1,1)),
            self._block(in_channels=32, out_channels=32, kernel_size=3, stride=(2,1),padding=(1,1)),
            self._block(in_channels=32, out_channels=1, kernel_size=3, stride=(1,1),padding=(1,1)),
            nn.Sigmoid()
            )
        
    
    def _block(self, in_channels, out_channels,kernel_size, stride,padding):
        
        return nn.Sequential(
            nn.ConvTranspose2d(in_channels,
                               out_channels,
                               kernel_size,
                               stride,
                               padding,
                               bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(), #they use relu in the generator
            )
    
    def forward(self, x):
        
        x = x.view(-1, 128)
        x = self.lin1(x)
        x = x.view(-1,256,6,4)
        
        return self.gen(x)
</code></pre>
<p>The inputs (real/fake/) have shape <code>(batch_size, 1, 85, 8)</code> and consist of very sparse one-hot matrices.</p>
<p>Now, with the above models, during the first training batches I have very bad errors for both <code>loss G</code> and <code>loss D</code></p>
<pre><code class="lang-auto">Epoch [0/5] Batch 0/84               Loss D: -34.0230, loss G: 132.8942
Epoch [0/5] Batch 1/84               Loss D: -3080.0264, loss G: 601.3990
Epoch [0/5] Batch 2/84               Loss D: -216907.8125, loss G: 872.5948
Epoch [0/5] Batch 3/84               Loss D: -26314.8633, loss G: 4973.5327
Epoch [0/5] Batch 4/84               Loss D: -1000911.5000, loss G: 6153.7974
Epoch [0/5] Batch 5/84               Loss D: -14484664.0000, loss G: -5013.7808
Epoch [0/5] Batch 6/84               Loss D: -5119665.0000, loss G: -7194.0640
Epoch [0/5] Batch 7/84               Loss D: -25285320.0000, loss G: 20130.0801
Epoch [0/5] Batch 8/84               Loss D: -11411679.0000, loss G: 32655.1016
Epoch [0/5] Batch 9/84               Loss D: -18403266.0000, loss G: 37912.0469
Epoch [0/5] Batch 10/84               Loss D: -6191229.0000, loss G: 33614.3828
Epoch [0/5] Batch 11/84               Loss D: -8119311.0000, loss G: 28472.3496
Epoch [0/5] Batch 12/84               Loss D: -134419216.0000, loss G: 18065.1074
Epoch [0/5] Batch 13/84               Loss D: -123661928.0000, loss G: 71028.8984
Epoch [0/5] Batch 14/84               Loss D: -2723217.0000, loss G: 47931.0195
Epoch [0/5] Batch 15/84               Loss D: -806806.1250, loss G: 41759.3555
</code></pre>
<p>Even though these are just the first batches of the first epoch, the losses seem too large to me and I suspect something’s wrong with my implementation. Or can be normal to obtain such numbers for the WGAN losses at first batches?</p>
<p>If the models look OK, should I upload my training loop for further discussion?</p>
<p>EDIT: I’m adding my training loop as it might help to figure out what’s happening here</p>
<pre><code class="lang-auto">opt_gen = optim.Adam(gen.parameters(), lr=0.001)

opt_critic = optim.Adam(critic.parameters(), lr = 0.0001)

# fixed_noise = torch.randn(32, Z_DIM, 1,1)

step=0

gen.train()

critic.train()

for epoch in range(N_EPOCHS):
    
    for batch_idx,real in enumerate(loader):
        
        
        #Maximizing the distance between the two probabilities p_G and p_data
        #TRAIN DISCRIMINATOR max (log(D(x))) + 1-log(D(G(z)))
        
        
        for _ in range(CRITIC_ITERATIONS):
            
            noise = torch.randn(real.shape[0], Z_DIM,1,1)
            fake = gen(noise)
            
            critic_real = critic(real).reshape(-1) 
            critic_fake = critic(fake).reshape(-1)
            
            gp = gradient_penalty(critic, real, fake, device='cpu')
            
            #we want to maximize here but algorithms like RMSProp are made for minimizing.
            # so we just use the trick of putting an extra minus sign.
            
            loss_critic = torch.mean(critic_fake) - torch.mean(critic_real) + LAMBDA_GP*gp
            
            critic.zero_grad()
            
            #retain graph cause we'll use fake for update step of generator.
            
            loss_critic.backward(retain_graph=True)
            
            opt_critic.step()
                
        
        # TRAIN GENERATOR: min -E[critic(gen_fake)]
        
        output = critic(fake).reshape(-1)
        
        loss_gen = -torch.mean(output)
        
        gen.zero_grad()
        
        loss_gen.backward()
        
        opt_gen.step()
        
        
        # Print losses and print to tensorboard
        
        print(
            f"Epoch [{epoch}/{N_EPOCHS}] Batch {batch_idx}/{len(loader)} \
              Loss D: {loss_critic:.4f}, loss G: {loss_gen:.4f}"
        )

        # with torch.no_grad():
        #     fake = gen(fixed_noise)
        #     # take out (up to) 32 examples
        #     img_grid_real = torchvision.utils.make_grid(
        #         real[:32], normalize=True
        #     )
        #     img_grid_fake = torchvision.utils.make_grid(
        #         fake[:32], normalize=True
        #     )

        #     writer_real.add_image("Real", img_grid_real, global_step=step)
        #     writer_fake.add_image("Fake", img_grid_fake, global_step=step)

        step += 1
</code></pre>