I was reading Improved Training of Wasserstein GANs 108, and thinking how it could be implemented in PyTorch. It seems not so complex but how to handle gradient penalty in loss troubles me.

2017-04-05.png709×125 6.71 KB

In the tensorflow’s implementation, the author use tf.gradients.

  
      github.com
  
  
    igul222/improved_wgan_training/blob/master/gan_cifar.py#L132-L136 98

gradients = tf.gradients(Discriminator(interpolates), [interpolates])[0]
slopes = tf.sqrt(tf.reduce_sum(tf.square(gradients), reduction_indices=[1]))
gradient_penalty = tf.reduce_mean((slopes-1.)**2)
disc_cost += LAMBDA*gradient_penalty





  
  
    
    
  
  


I wonder if there is an easy way to handle the gradient penalty.
here is my idea of implementing it, I don’t know whether it will work and work in the way I think:

optimizer_D = optim.adam(model_D.parameters())
x = Variable()
y = Variable()
x_hat = (alpha*x+(1-alpha)*y).detach()
x_hat.requires_grad = True

loss_D = model_D(x_hat).sum()
loss_D.backward()
x_hat.grad.volatile = False

loss = model_D(x).sum() - model_D(y).sum() + ((x_hat.grad -1)**2 * LAMBDA).sum()
loss.backward()
optimizer_D.step()