Don't keep unnecessary saved_inputs alive (#16583)
authorAdam Paszke <adam.paszke@gmail.com>
Mon, 11 Feb 2019 21:31:06 +0000 (13:31 -0800)
committerFacebook Github Bot <facebook-github-bot@users.noreply.github.com>
Mon, 11 Feb 2019 21:42:09 +0000 (13:42 -0800)
commit7743ed850231781097fe4388fbe6d299cbb4a2c6
treea05bf38b8fd81cf3f42c156557b3736b3ed318d2
parente2a5b203fccf78717c5171e89e2f9d258c487fe2
Don't keep unnecessary saved_inputs alive (#16583)

Summary:
Fixes #16577.

This greatly improves memory efficiency of certain ops like Dropout2d. Previously, they were implemented as `input * mask` where mask never requires_grad, but we didn't use that knowledge in forward, and (in case of a in-place dropout) kept input.clone() for the backward, when it would simply get ignored.

This patch tries to address this situation by emitting some guards for stores like this, but only if they are as simple, as checking if a single value requires_grad.

Interestingly, the same optimizations apply to methods like bmm, baddmm, etc., but _not to mm nor addmm_, because of how their derivatives are defined. Apparently they unnecessarily use `mat1` to compute the derivative of `mat1` just to improve the error message in case `mat1` was sparse. I'd like to apply this optimization to that case, but I don't want to loose the nicer error message, so if anyone has any ideas for solutions, please let me know...

Full list of operators affected by this patch:
* _nnpack_spatial_convolution
* addbmm
* addcdiv
* addcmul
* addmv
* addr
* baddbmm
* bmm
* cross
* div
* dot
* fmod
* ger
* index_add_
* mul
* mv
* scatter_add_
Pull Request resolved: https://github.com/pytorch/pytorch/pull/16583

Differential Revision: D13900881

Pulled By: gchanan

fbshipit-source-id: dd0aeb2ab58c4b6aa95b37b46d3255b3e014291c
test/test_autograd.py
tools/autograd/gen_variable_type.py