Pytorch how to get the gradient of loss function twice
Here is what I'm trying to implement:

We calculate loss based on F(X), as usual. But we also define "adversarial loss" which is a loss based on F(X + e). e is defined as dF(X)/dX multiplied by some constant. Both loss and adversarial loss are backpropagated for the total loss.

In tensorflow, this part (getting dF(X)/dX) can be coded like below:

  grad, = tf.gradients( loss, X )
  grad = tf.stop_gradient(grad)
  e = constant * grad


Below is my pytorch code:

class DocReaderModel(object):
    def __init__(self, embedding=None, state_dict=None):
        self.train_loss = AverageMeter()
        self.embedding = embedding
        self.network = DNetwork(opt, embedding)
        self.optimizer = optim.SGD(parameters)

    def adversarial_loss(self, batch, loss, embedding, y):
        self.optimizer.zero_grad()
        loss.backward(retain_graph=True)
        grad = embedding.grad
        grad.detach_()

        perturb = F.normalize(grad, p=2)* 0.5
        self.optimizer.zero_grad()
        adv_embedding = embedding + perturb
        network_temp = DNetwork(self.opt, adv_embedding) # This is how to get F(X)
        network_temp.training = False
        network_temp.cuda()
        start, end, _ = network_temp(batch) # This is how to get F(X)
        del network_temp # I even deleted this instance.
        return F.cross_entropy(start, y[0]) + F.cross_entropy(end, y[1])

    def update(self, batch):
        self.network.train()
        start, end, pred = self.network(batch)
        loss = F.cross_entropy(start, y[0]) + F.cross_entropy(end, y[1])
        loss_adv = self.adversarial_loss(batch, loss, self.network.lexicon_encoder.embedding.weight, y) 
        loss_total = loss + loss_adv 

        self.optimizer.zero_grad()
        loss_total.backward()
        self.optimizer.step()


I have few questions:

1) I substituted tf.stop_gradient with grad.detach_(). Is this correct?

2) I was getting "RuntimeError: Trying to backward through the graph a second time, but the buffers have already been freed. Specify retain_graph=True when calling backward the first time." so I added retain_graph=True at the loss.backward. That specific error went away.
However now I'm getting a memory error after few epochs (RuntimeError: cuda runtime error (2) : out of memory at /opt/conda/conda-bld/pytorch_1525909934016/work/aten/src/THC/generic/THCStorage.cu:58
). I suspect I'm unnecessarily retaining graph. 

Can someone let me know pytorch's best practice on this? Any hint / even short comment will be highly appreciated.

I think you are trying to implement generative adversarial network (GAN), but from the code, I don't understand and can't follow to what you are trying to achieve as there are a few missing pieces for a GAN to works. I can see there's a discriminator network module, DNetwork but missing the generator network module.

If to guess, when you say 'loss function twice', I assumed you mean you have one loss function for the discriminator net and another for the generator net. If that's the case, let me share how I would implement a basic GAN model.

As an example, let's take a look at this Wasserstein GAN Jupyter notebook

I'll skip the less important bits and zoom into the important ones here:


First, import PyTorch libraries and set up

# Set up batch size, image size, and size of noise vector:
bs, sz, nz = 64, 64, 100 # nz is the size of the latent z vector for creating some random noise later

Build a discriminator module

class DCGAN_D(nn.Module):
    def __init__(self):
        ... truncated, the usual neural nets stuffs, layers, etc ...
    def forward(self, input):
        ... truncated, the usual neural nets stuffs, layers, etc ...

Build a generator module

class DCGAN_G(nn.Module):
    def __init__(self):
        ... truncated, the usual neural nets stuffs, layers, etc ...
    def forward(self, input):
        ... truncated, the usual neural nets stuffs, layers, etc ...

Put them all together

netG = DCGAN_G().cuda()
netD = DCGAN_D().cuda()

Optimizer needs to be told what variables to optimize. A module automatically keeps track of its variables.

optimizerD = optim.RMSprop(netD.parameters(), lr = 1e-4)
optimizerG = optim.RMSprop(netG.parameters(), lr = 1e-4)

One forward step and one backward step for Discriminator

Here, the network can calculate gradient during the backward pass, depends on the input to this function. So, in my case, I have 3 type of losses; generator loss, dicriminator real image loss, dicriminator fake image loss. I can get gradient of loss function three times for 3 different net passes.

def step_D(input, init_grad):
    # input can be from generator's generated image data or input image from dataset
    err = netD(input)
    err.backward(init_grad) # backward pass net to calculate gradient
    return err # loss

Control trainable parameters [IMPORTANT]

Trainable parameters in the model are those that require gradients.

def make_trainable(net, val):
    for p in net.parameters():
        p.requires_grad = val # note, i.e, this is later set to False below in netG update in the train loop.


In TensorFlow, this part can be coded like below:

grad = tf.gradients(loss, X)
grad = tf.stop_gradient(grad)


So, I think this will answer your first question, "I substituted tf.stop_gradient with grad.detach_(). Is this correct?"
Train loop


You can see here how's the 3 different loss functions are being called here.

    def train(niter, first=True):

        for epoch in range(niter):
            # Make iterable from PyTorch DataLoader
            data_iter = iter(dataloader)
            i = 0

            while i &lt; n:
                ###########################
                # (1) Update D network
                ###########################
                make_trainable(netD, True)

                # train the discriminator d_iters times
                d_iters = 100

                j = 0

                while j &lt; d_iters and i &lt; n:
                    j += 1
                    i += 1

                    # clamp parameters to a cube
                    for p in netD.parameters():
                        p.data.clamp_(-0.01, 0.01)

                    data = next(data_iter)

                    ##### train with real #####
                    real_cpu, _ = data
                    real_cpu = real_cpu.cuda()
                    real = Variable( data[0].cuda() )
                    netD.zero_grad()

                    # Real image discriminator loss
                    errD_real = step_D(real, one)

                    ##### train with fake #####
                    fake = netG(create_noise(real.size()[0]))
                    input.data.resize_(real.size()).copy_(fake.data)

                    # Fake image discriminator loss
                    errD_fake = step_D(input, mone)

                    # Discriminator loss
                    errD = errD_real - errD_fake
                    optimizerD.step()

                ###########################
                # (2) Update G network
                ###########################
                make_trainable(netD, False)
                netG.zero_grad()

                # Generator loss
                errG = step_D(netG(create_noise(bs)), one)
                optimizerG.step()

                print('[%d/%d][%d/%d] Loss_D: %f Loss_G: %f Loss_D_real: %f Loss_D_fake %f'
                    % (epoch, niter, i, n,
                    errD.data[0], errG.data[0], errD_real.data[0], errD_fake.data[0]))



  "I was getting "RuntimeError: Trying to backward through the graph a second time..."


PyTorch has this behaviour; to reduce GPU memory usage, during the .backward() call, all the intermediary results (if you have like saved activations, etc.) are deleted when they are not needed anymore. Therefore, if you try to call .backward() again, the intermediary results don't exist and the backward pass cannot be performed (and you get the error you see).

It depends on what you are trying to do. You can call .backward(retain_graph=True) to make a backward pass that will not delete intermediary results, and so you will be able to call .backward() again. All but the last call to backward should have the retain_graph=True option.


  Can someone let me know pytorch's best practice on this


As you can see from the PyTorch code above and from the way things are being done in PyTorch which is trying to stay Pythonic, you can get a sense of PyTorch's best practice there.
