How to correctly apply spectral normalization for WGAN?

ZimoNitrome: <p>I am using <a href="https://github.com/aladdinpersson" rel="noopener nofollow ugc">aladdinpersson</a>s code for WGAN-GP: <a href="https://github.com/aladdinpersson/Machine-Learning-Collection/blob/master/ML/Pytorch/GANs/4.%20WGAN-GP/train.py" class="inline-onebox" rel="noopener nofollow ugc">Machine-Learning-Collection/ML/Pytorch/GANs/4. WGAN-GP/train.py at master · aladdinpersson/Machine-Learning-Collection · GitHub</a></p>
<p>and removed the GP part and instead applied spectral normalization to both the critic and the generator as follows:</p>
<pre data-code-wrap="python"><code class="lang-python">critic_fake = critic(fake).reshape(-1)
gp = gradient_penalty(critic, real, fake, device=device)
loss_critic = (
    -(torch.mean(critic_real) - torch.mean(critic_fake)) + LAMBDA_GP * gp
)
critic.zero_grad()
</code></pre>
<p>becomes:</p>
<pre data-code-wrap="python"><code class="lang-python">critic_fake = critic(fake).reshape(-1)
loss_critic = (
    -(torch.mean(critic_real) - torch.mean(critic_fake))
)
critic.zero_grad()
</code></pre>
<p>and:</p>
<pre data-code-wrap="python"><code class="lang-python"># initialize gen and disc, note: discriminator should be called critic,
# according to WGAN paper (since it no longer outputs between [0, 1])
gen = Generator(Z_DIM, CHANNELS_IMG, FEATURES_GEN).to(device)
critic = Discriminator(CHANNELS_IMG, FEATURES_CRITIC).to(device)
initialize_weights(gen)
initialize_weights(critic)
</code></pre>
<p>becomes:</p>
<pre data-code-wrap="python"><code class="lang-python"># initialize gen and disc, note: discriminator should be called critic,
# according to WGAN paper (since it no longer outputs between [0, 1])
gen = Generator(Z_DIM, CHANNELS_IMG, FEATURES_GEN).to(device)
critic = Discriminator(CHANNELS_IMG, FEATURES_CRITIC).to(device)
initialize_weights(gen)
initialize_weights(critic)
apply_spectral_norm(gen)
apply_spectral_norm(critic)
</code></pre>
<p>Where <code>apply_spectral_norm()</code> is defined as:</p>
<pre data-code-wrap="python"><code class="lang-python">def flatten_network(net):
    flattened = [flatten_network(children) for children in net.children()]
    res = [net]
    for c in flattened:
        res += c
    return res

def apply_spectral_norm(net):
    # Apply spectral normalization for conv layers
    for p in flatten_network(net):
        if isinstance(p, (nn.Conv2d, nn.ConvTranspose2d, nn.Linear)):
            nn.utils.spectral_norm(p)
</code></pre>
<p>I expected the GAN to converge in similar speed but this is what I see:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/1/f/1fac521e788fc5039bfce9eb1bf7d9b0d50c710f.png" data-download-href="https://discuss.pytorch.org/uploads/default/1fac521e788fc5039bfce9eb1bf7d9b0d50c710f" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/1/f/1fac521e788fc5039bfce9eb1bf7d9b0d50c710f_2_276x500.png" alt="image" data-base62-sha1="4wbZP21yNPRbPKvo3eck6kVqel1" width="276" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/1/f/1fac521e788fc5039bfce9eb1bf7d9b0d50c710f_2_276x500.png, https://discuss.pytorch.org/uploads/default/original/3X/1/f/1fac521e788fc5039bfce9eb1bf7d9b0d50c710f.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/1/f/1fac521e788fc5039bfce9eb1bf7d9b0d50c710f.png 2x" data-dominant-color="F9F9FA"><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">395×715 21 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>(blue is WGAN-GP and red is WGAN-SN)</p>
<p>Any tips on what I am missing?</p>

ZimoNitrome: <p>bumping this because I still haven’t found a good configuration</p>

ZimoNitrome: <p>Maaaybe I have found 2 improvements that work.</p>
<p><span class="hashtag">#1:</span> I think possibly my way of applying spectral norm was incorrect…</p>
<p>Now, I try to apply it to all conv, linear, and normalization layers:</p>
<pre><code class="lang-python">def apply_spectral_norm(net: torch.nn.Module, power=1):
    for m in net.modules():
        if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d, nn.Linear, nn.BatchNorm2d, nn.BatchNorm1d, nn.InstanceNorm2d, nn.InstanceNorm1d)):
            nn.utils.parametrizations.spectral_norm(m, n_power_iterations=power)
</code></pre>
<p>This also uses the updated and non-deprecated way of adding spectral norm.</p>
<p>Maybe there is a better way to do this? Like check if “weight” exists in each module somehow.</p>
<p><span class="hashtag">#2:</span> Now I only apply it to the discriminator:</p>
<pre><code class="lang-python"># initialize gen and disc, note: discriminator should be called critic,
# according to WGAN paper (since it no longer outputs between [0, 1])
gen = Generator(Z_DIM, CHANNELS_IMG, FEATURES_GEN).to(device)
critic = Discriminator(CHANNELS_IMG, FEATURES_CRITIC).to(device)
initialize_weights(gen)
initialize_weights(critic)
# apply_spectral_norm(gen)
apply_spectral_norm(critic)
</code></pre>
<p>This seems to work better than previously.</p>

ZimoNitrome: <p>But trying this out on the minimal example stated above, it still does not compare at all:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/e/e/eeb0802c44468eca9bd51a5b49145ddeb036073a.png" data-download-href="https://discuss.pytorch.org/uploads/default/eeb0802c44468eca9bd51a5b49145ddeb036073a" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/e/e/eeb0802c44468eca9bd51a5b49145ddeb036073a_2_480x500.png" alt="image" data-base62-sha1="y3xMUBJnLTRhlplj6CAHxxQBAZI" width="480" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/e/e/eeb0802c44468eca9bd51a5b49145ddeb036073a_2_480x500.png, https://discuss.pytorch.org/uploads/default/optimized/3X/e/e/eeb0802c44468eca9bd51a5b49145ddeb036073a_2_720x750.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/e/e/eeb0802c44468eca9bd51a5b49145ddeb036073a_2_960x1000.png 2x" data-dominant-color="F5F5F5"><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">1104×1150 66.9 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><br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/1/4/14225f2f12cdff4f550f44815f37ec1115b558ad.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/14225f2f12cdff4f550f44815f37ec1115b558ad" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/1/4/14225f2f12cdff4f550f44815f37ec1115b558ad_2_690x250.jpeg" alt="image" data-base62-sha1="2S7bOXA0AsJCYTjvfXY06McWQCV" width="690" height="250" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/1/4/14225f2f12cdff4f550f44815f37ec1115b558ad_2_690x250.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/1/4/14225f2f12cdff4f550f44815f37ec1115b558ad_2_1035x375.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/1/4/14225f2f12cdff4f550f44815f37ec1115b558ad_2_1380x500.jpeg 2x" data-dominant-color="777676"><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">1702×617 169 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>

hidden_machine: <p>If I’m not wrong, you’re supposed to apply spectral norm every iteration, and typically to only the critic. Also, increase the value of n_power_iterations if you haven’t, because it gives a better estimate of the spectral norm.</p>
<aside class="quote no-group" data-username="ZimoNitrome" data-post="3" data-topic="137427">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/zimonitrome/48/27591_2.png" class="avatar"> ZimoNitrome:</div>
<blockquote>
<p>Maybe there is a better way to do this? Like check if “weight” exists in each module somehow.</p>
</blockquote>
</aside>
<p>Typically the learnable layers like convolution layers and dense layers have weights in the critic, you can manually apply it to that instead of having a function to check. I don’t think you need to have any other type of layer in the critic network afaik, cause things like batch normalization don’t help much in the critic network.</p>
<aside class="quote no-group" data-username="ZimoNitrome" data-post="4" data-topic="137427">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/zimonitrome/48/27591_2.png" class="avatar"> ZimoNitrome:</div>
<blockquote>
<p>But trying this out on the minimal example stated above, it still does not compare at all:</p>
</blockquote>
</aside>
<p>Did you try running it for more time? My guess is that since you are dividing by the spectral norm, the weights become closer to zero → so smaller gradients are being backpropagated → so the learning in the previous layers and the generator are slower. Increase the learning rate in case of the spectral norm example or just wait a bit longer.</p>