Performance improvement

[htt210]

I found that performing forward-backward in a loop is much faster than using autograd.grad with retain_graph=True. Your current code is:

        loglikelihood_grads = zip(*[autograd.grad(
            l, self.parameters(),
            retain_graph=(i < len(loglikelihoods))
        ) for i, l in enumerate(loglikelihoods, 1)])

https://github.com/kuc2477/pytorch-ewc/blob/master/model.py#L75

after I change it to:

        for batch, label in zip(buffer_data, buffer_label):
            self.zero_grad()
            loglikelihood = F.cross_entropy(self(batch), label)
            loglikelihood.backward()
            for n, p in self.named_parameters():
                n = n.replace('.', '__')
                grads[n] = grads.get(n, 0) + p.grad ** 2

it runs much faster. Can you please investigate this?

Thank you.
Thanh Tung