The aixs of norm of the gradient

[CharlesNord]

wgan-gp/gan_mnist.py

Line 148 in ae47a18

gradient_penalty = ((gradients.norm(2, dim=1) - 1) ** 2).mean() * LAMBDA

Hi, I accidently saw your code when I google W-GAN GP. I think there is something wrong with your implementationi here. In W-GAN GP, the norm of the interpolated gradient should be calculated across all axis except the batch axis, since the gradient is wrt each sample. But in your code, you only calculated the norm of the second dimension, which is not reaonable. I think you miss the following reshape step:

gradients.view(gradients.shape[0], -1)

[1292224662]

I also noticed this mistake, and instead of using torch.view(), I used the following:

 gradient_penalty = ((gradients.norm(2, dim=(1,2,3)) - 1) ** 2).mean() * LAMBDA 

Since the gradients here is in the shape of [batch_size, 3, H, W], I think this can also get the L2 norm of each batches, right?