Truncating backprop temporarily

Kaixhin · July 18, 2017, 12:03am

Is it possible to accomplish the following (from Chainer) in PyTorch?

chainer/chainerrl/blob/master/chainerrl/agents/acer.py#L193-L200


@contextlib.contextmanager
def backprop_truncated(*variables):
backup = [v.creator for v in variables]
for v in variables:
    v.unchain()
yield
for v, backup_creator in zip(variables, backup):
    v.set_creator(backup_creator)

For reference, I’m trying to implement the “efficient” trust region optimisation from ACER: Trust region update. Whereas my code currently backprops through the entire computation graph of the policy gradient loss and is hence very expensive, the efficient version cuts the graph just before the final layer, then calculates gradients. However, the parents need to be rejoined, as eventually the new loss should be backpropped through the entire graph.

tom · July 18, 2017, 3:41am

You could split it by default in the forward pass and manually pass the gradients of the last bit to the backward of what has been cut off when that is desired.

Best regards

Thomas