PLEASE give me my memory back

1

Pytorch refuses to release the computation graph no matter what I try. I am trying to incorporate the loss of the gradient with respect to the input, but this is seemingly unprecedented in the Pytorch community, has it never been done before?!

Relevant code:

    def train(self, train_loader, epoch, epochs):
        self.net.train()
        train_loss = 0.0
        for i, data in enumerate(train_loader):
            features, labels, label_derivatives = data

            features.requires_grad = True
            outputs = self.net(features)
            outputs.backward(features, create_graph=True)
            derivatives = features.grad

            self.optimizer.zero_grad()
            loss1 = (outputs - labels).pow(2).mean()
            loss2 = (derivatives - label_derivatives).pow(2).mean()
            loss = loss1 + loss2
            loss.backward()
            self.optimizer.step()
            train_loss += float(loss.item())

        return train_loss / len(train_loader)
2

If this is not about the cached memory, use torch.autograd.grad instead of outputs.backward(xxxx, ...). See https://github.com/pytorch/pytorch/issues/7343 for details.

3

Yeah, I fixed it by doing

features.requires_grad = True
outputs = self.net(features)
derivatives = torch.autograd.grad(
    outputs=outputs,
    inputs=features,
    grad_outputs=torch.ones_like(outputs),
    create_graph=True
)[0]

In case anyone runs into the same problem.