A mistake on gradient penalty

[doub7e]

https://github.com/Zeleni9/pytorch-wgan/blob/master/models/wgan_gradient_penalty.py#L324

grad_penalty = ((gradients.norm(2, dim=1) - 1) ** 2).mean() * self.lambda_term

The shape of gradients is [batch_size, 1, 32, 32], the above norm(2, dim=1) is actually not behaving as you may want (and returns a tensor with shape [batch_size, 32, 32]). This is actually enforcing gradients to be closer to 1 element-wise. So the gradient and the lipschitz are much bigger than 1.

A potential fix is
grad_penalty = (((gradients.view(gradients.shape[0], -1) ** 2).sum(dim=1).sqrt() - 1) ** 2).mean() * self.lambda_term

Correct me if I am wrong. Thx.

[R-N]

You're right. The gradient penalty is now stable and the generator loss isn't changing so fast. I prefer to just use tensor.norm though.

grad_norm = gradients.view(gradients.shape[0], -1).norm(2, dim=-1)
gradient_penalty = ((grad_norm - 1) ** 2).mean()

[R-N]

Ok, well, this adjustment actually adds artifacts to the generated images? Not sure why

With adjustment:

Without:

[R-N]

Tried your code, and yeah, it has artifacts too. I wonder what's wrong.
Perhaps this is one of those "if it's not broken don't fix it" moments.

[R-N]

Took a look at the official WGAN GP code and it doesn't have batch-norm. So I removed it and the artifacts are gone. It doesn't seem to be better than before the adjustment though.

EDIT: Actually batch norm is only not used for MNIST, but unless I disable it for others, I'm still getting artifacts.