SRGAN artifacts caused by gradient

Valentin_Fontanger: <p>Hi,</p>
<p>I am trying to implement SRGAN (super resolution).<br>
I try to train my model on a batch of 10 images, just in order to see if I am able to correctly overfit.<br>
Here are the results after 200 epochs on these 10 images :</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/1/3/13616b9040f8f378885d0049b20f875232a66dde.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/13616b9040f8f378885d0049b20f875232a66dde" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/1/3/13616b9040f8f378885d0049b20f875232a66dde_2_690x269.jpeg" alt="image" data-base62-sha1="2LrNcWwGOmmm9ky48vO4kWWc8rQ" width="690" height="269" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/1/3/13616b9040f8f378885d0049b20f875232a66dde_2_690x269.jpeg, https://discuss.pytorch.org/uploads/default/original/3X/1/3/13616b9040f8f378885d0049b20f875232a66dde.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/1/3/13616b9040f8f378885d0049b20f875232a66dde.jpeg 2x" data-dominant-color="2C2A2A"><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">image</span><span class="informations">887×346 71.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>I highly suspect my training loop, and my non understanding of how the gradient flows with pytorch.<br>
Here is the code :</p>
<pre><code class="lang-auto">torch.autograd.set_detect_anomaly(True)

for epoch in range(epochs): 
    
    for i, batch in enumerate(train_data): 
        
        lr, hr = batch["lr"], batch["hr"]
        lr = lr.to(device)
        hr = hr.to(device)
        bs = lr.shape[0]
        # train discriminator
        
        
        label = torch.full((bs, ), 1, device=device, dtype=torch.float32)
        d_net.zero_grad()
        real_pred = d_net(hr).view(-1)
                
        real_loss = bce_criterion(real_pred,label)
        real_loss.backward()
        # change label 
        label = label.fill_(0.0)
        sr = g_net(lr)
        fake_pred = d_net(sr.detach()).view(-1)
        fake_loss = bce_criterion(fake_pred, label)
        fake_loss.backward()
        d_loss = fake_loss + real_loss
        d_optim.step()
        
        with torch.no_grad():
            p_real = real_pred.mean(axis=0)
            p_fake = fake_pred.mean(axis=0)
        
        # train generator
        
        # change label 
        g_net.zero_grad()
        label = label.fill_(1.)
        sr_preds = d_net(sr).view(-1)
        adv_loss = adv_coeff * bce_criterion(sr_preds, label)
        adv_loss.backward(retain_graph=True)
        content_criterion = content_loss(hr, sr)
        content_criterion.backward()
        g_loss = content_criterion + adv_loss
        g_optim.step()
        
        with torch.no_grad():
            p_fake_trick = sr_preds.mean(axis=0)

    
        # monitoring
                
        if i % 10 == 0: 
            print("d_loss = ", d_loss.item())
            print("g_loss = ", g_loss.item())
            print("D(Real) = ", p_real.item())
            print("D(Fake)2 = ", p_fake_trick.item())
            print("Content loss = ", content_criterion.item())

            with torch.no_grad(): 
                for test_batch in train_data: 
                    
                    lr_test, hr_test = test_batch["lr"], test_batch["hr"]
                    lr_test = lr_test.to(device)
                    hr_test = hr_test.to(device)
                    print("LR")
                    display_grid(lr_test.cpu())
                    sr = g_net(lr_test)
                    print("Generated")
                    display_grid(sr.cpu())
                    print("Real")
                    display_grid(hr_test.cpu())
                    break
            
    print("Epoch : ", epoch)
</code></pre>
<p>Thank you very very much for your help</p>