Fix memory leak in Adam, Adagrad, RMSProp (#23125)

jerry73204 · facebook-github-bot · commit 2ac9abf75918 · 2019-07-21T10:06:18.000-07:00
Summary: As reported in LaurentMazare/tch-rs#76, the memory grows when weight_decay is present when using Adam. It applies the same fix in #23007 to Adam, Adagrad and RMSProp. Pull Request resolved: #23125 Differential Revision: D16402421 Pulled By: soumith fbshipit-source-id: 59eb4bd81b8bd9e1a5f7c068ed841f70a4c38a80
diff --git a/torch/csrc/api/src/optim/adagrad.cpp b/torch/csrc/api/src/optim/adagrad.cpp
@@ -24,6 +24,7 @@ void Adagrad::step() {
     }
 
     if (options.weight_decay_ > 0) {
+      NoGradGuard guard;
       p.grad() = p.grad() + options.weight_decay_ * p;
     }
 
diff --git a/torch/csrc/api/src/optim/adam.cpp b/torch/csrc/api/src/optim/adam.cpp
@@ -23,6 +23,7 @@ void Adam::step() {
     }
 
     if (options.weight_decay_ > 0) {
+      NoGradGuard guard;
       p.grad() = p.grad() + options.weight_decay_ * p;
     }
 
diff --git a/torch/csrc/api/src/optim/rmsprop.cpp b/torch/csrc/api/src/optim/rmsprop.cpp
@@ -24,6 +24,7 @@ void RMSprop::step() {
     }
 
     if (options.weight_decay_ > 0) {
+      NoGradGuard guard;
       p.grad() = p.grad() + options.weight_decay_ * p;
     }
 

Original file line number	Diff line number	Diff line change
`@@ -24,6 +24,7 @@ void Adagrad::step() {`
`24`	`24`	`}`
`25`	`25`
`26`	`26`	`if (options.weight_decay_ > 0) {`
	`27`	`+ NoGradGuard guard;`
`27`	`28`	`p.grad() = p.grad() + options.weight_decay_ * p;`
`28`	`29`	`}`
`29`	`30`
Original file line number	Diff line number	Diff line change
`@@ -23,6 +23,7 @@ void Adam::step() {`
`23`	`23`	`}`
`24`	`24`
`25`	`25`	`if (options.weight_decay_ > 0) {`
	`26`	`+ NoGradGuard guard;`
`26`	`27`	`p.grad() = p.grad() + options.weight_decay_ * p;`
`27`	`28`	`}`
`28`	`29`
Original file line number	Diff line number	Diff line change
`@@ -24,6 +24,7 @@ void RMSprop::step() {`
`24`	`24`	`}`
`25`	`25`
`26`	`26`	`if (options.weight_decay_ > 0) {`
	`27`	`+ NoGradGuard guard;`
`27`	`28`	`p.grad() = p.grad() + options.weight_decay_ * p;`
`28`	`29`	`}`
`29`	`30`