Fix refcycles in DataParallel scatter and gather (#4988)

zou3519 · soumith · commit da3c4cb84c27 · 2018-02-07T09:17:34.000-08:00
* Eliminate reference cycles in scatter_gather

* Test for refcycles

* Better fix

* Add comments
diff --git a/torch/nn/parallel/scatter_gather.py b/torch/nn/parallel/scatter_gather.py
@@ -22,7 +22,15 @@ def scatter_map(obj):
             return list(map(type(obj), zip(*map(scatter_map, obj.items()))))
         return [obj for targets in target_gpus]
 
-    return scatter_map(inputs)
+    # After scatter_map is called, a scatter_map cell will exist. This cell
+    # has a reference to the actual function scatter_map, which has references
+    # to a closure that has a reference to the scatter_map cell (because the
+    # fn is recursive). To avoid this reference cycle, we set the function to
+    # None, clearing the cell
+    try:
+        return scatter_map(inputs)
+    finally:
+        scatter_map = None
 
 
 def scatter_kwargs(inputs, kwargs, target_gpus, dim=0):
@@ -50,4 +58,10 @@ def gather_map(outputs):
         if out is None:
             return None
         return type(out)(map(gather_map, zip(*outputs)))
-    return gather_map(outputs)
+
+    # Recursive function calls like this create reference cycles.
+    # Setting the function to None clears the refcycle.
+    try:
+        return gather_map(outputs)
+    finally:
+        gather_map = None
