I am trying to implement a loss function in Pytorch which requires me to get the gradient of the logit of a discriminator of a GAN as such: 
image.jpg779×72 12.4 KB

There is an implementation of this in tensorflow already which uses the tf.gradients function:
   # -----------------------------------------------------------------------------------
    #     JS-Regularizer
    # -----------------------------------------------------------------------------------
    def Discriminator_Regularizer(D1_logits, D1_arg, D2_logits, D2_arg):
        with tf.name_scope('disc_reg'):
            D1 = tf.nn.sigmoid(D1_logits)
            D2 = tf.nn.sigmoid(D2_logits)
            grad_D1_logits = tf.gradients(D1_logits, D1_arg)[0]
            grad_D2_logits = tf.gradients(D2_logits, D2_arg)[0]
            grad_D1_logits_norm = tf.norm( tf.reshape(grad_D1_logits, [BATCH_SIZE//len(DEVICES),-1]) , axis=1)
            grad_D2_logits_norm = tf.norm( tf.reshape(grad_D2_logits, [BATCH_SIZE//len(DEVICES),-1]) , axis=1)
            
            #set keep_dims=True/False such that grad_D_logits_norm.shape == D.shape
            print('grad_D1_logits_norm.shape {} != D1.shape {}'.format(grad_D1_logits_norm.shape, D1.shape))
            print('grad_D2_logits_norm.shape {} != D2.shape {}'.format(grad_D2_logits_norm.shape, D2.shape))
            assert grad_D1_logits_norm.shape == D1.shape
            assert grad_D2_logits_norm.shape == D2.shape
            
            reg_D1 = tf.multiply(tf.square(1.0-D1), tf.square(grad_D1_logits_norm))
            reg_D2 = tf.multiply(tf.square(D2), tf.square(grad_D2_logits_norm))
            
            disc_regularizer = tf.reduce_mean(reg_D1 + reg_D2)
            
     return disc_regularizer

Does anyone know how to do the equivalent using autograd? Is this functionality possible in pytorch? It would be greatly appreciated.
My understanding is that tf.gradients(ys,xs) makes the symbolic derivatives of the sum of ys wrt to x. How is this different from autograd.grad()?
Michael