Computing gradients only for the front-end network in Pytorch

user_7386954 0
Viewed 47

I have a very simple question.

Let's say that I have two networks to train (i.e., net1, net2). The output of net1 will be fed into net2 while training. In my case, I would like to only update net1:

optimizer=Optimizer(net1.parameters(), **kwargs)
loss=net2(net1(x))
loss.backward()
optimizer.step()

Although this will achieve what I'm aiming for, it takes up too much redundant memory since this will compute the gradients for net2 (causes OOM error). Therefore I have tried out several attempts to solve this issue:

  1. torch.no_grad:
z=net1(x)
with torch.no_grad():
    loss=net2(z)

Didn't raise OOM but removed all the gradients including the ones from net1.

  1. requires_grad=False:
net2.requires_grad=False
loss=net2(net1(x))

Raised OOM.

  1. detach():
z=net1(x)
loss=net2(z).detach()

Didn't raise OOM but removed all the gradients including the ones from net1.

  1. eval():
net2.eval()
loss=net2(net1(x))

Raised OOM.

Is there any way to compute the gradients only for the front-end network (net1) for memory efficiency? Any suggestions would be appreciated.

pytorch
1 Answers

First let's try to understand why your methods don't work.

  1. This context manager disables all gradient computation.
  2. Since net1 requires a gradient, the subsequent requires_grad=False are ignored.
  3. If you detach at that state, this means the gradient computation already stops right there
  4. Eval just sets net2 to eval mode, that doesn't influence the gradient computation at all.

Depending on your architecture, the OOM error could already come from saving all the intermediate values in your computation graph (frequently an issue in CNNs), or it could come from having to store gradients (more common in fully connected networks).

What you probably are looking for is called "checkpointing" which you don't even have to implement on your own, you can use the checkpointing API of pytorch, check out the documentation.

This basically lets you compute and process the gradient for net1 and net2 separately. Note that you do need all the gradient informatino to go through net2, otherwise you cannot compute the gradients wrt. net1!

user_29131060
Related