Cleanup of the neural style transfer to reduce memory usage (#1665)

albanD · brianjo · web-flow · commit 067927177d90 · 2021-09-01T12:53:06.000-04:00
* No need to clone the model, it is used only once

* cleanup use of .data

* Set requires_grad properly to avoid computing un-needed gradients

Co-authored-by: Brian Johnson &lt;brianjo@fb.com&gt;
diff --git a/advanced_source/neural_style_tutorial.py b/advanced_source/neural_style_tutorial.py
@@ -309,8 +309,6 @@ def get_style_model_and_losses(cnn, normalization_mean, normalization_std,
                                style_img, content_img,
                                content_layers=content_layers_default,
                                style_layers=style_layers_default):
-    cnn = copy.deepcopy(cnn)
-
     # normalization module
     normalization = Normalization(normalization_mean, normalization_std).to(device)
 
@@ -394,7 +392,7 @@ def get_style_model_and_losses(cnn, normalization_mean, normalization_std,
 
 def get_input_optimizer(input_img):
     # this line to show that input is a parameter that requires a gradient
-    optimizer = optim.LBFGS([input_img.requires_grad_()])
+    optimizer = optim.LBFGS([input_img])
     return optimizer
 
 
@@ -418,6 +416,12 @@ def run_style_transfer(cnn, normalization_mean, normalization_std,
     print('Building the style transfer model..')
     model, style_losses, content_losses = get_style_model_and_losses(cnn,
         normalization_mean, normalization_std, style_img, content_img)
+
+    # We want to optimize the input and not the model parameters so we
+    # update all the requires_grad fields accordingly
+    input_img.requires_grad_(True)
+    model.requires_grad_(False)
+
     optimizer = get_input_optimizer(input_img)
 
     print('Optimizing..')
@@ -426,7 +430,8 @@ def run_style_transfer(cnn, normalization_mean, normalization_std,
 
         def closure():
             # correct the values of updated input image
-            input_img.data.clamp_(0, 1)
+            with torch.no_grad():
+                input_img.clamp_(0, 1)
 
             optimizer.zero_grad()
             model(input_img)
@@ -456,7 +461,8 @@ def closure():
         optimizer.step(closure)
 
     # a last correction...
-    input_img.data.clamp_(0, 1)
+    with torch.no_grad():
+        input_img.clamp_(0, 1)
 
     return input_img
 
