Repository: OpenMined/PySyft

Issue:
Garbage collection: memory leak in SMPC
**Describe the bug**
Garbage collection is not working as expected in SMPC

**To Reproduce**
```python
import torch as th
import syft as sy

hook = sy.TorchHook(th)
bob = sy.VirtualWorker(hook, id="bob")
alice = sy.VirtualWorker(hook, id="alice")
crypto_provider = sy.VirtualWorker(hook, id="james")

torch = th
syft = sy

class Classifier(torch.nn.Module):
    def __init__(self, in_features, out_features):
        super(Classifier, self).__init__()
        self.fc = torch.nn.Linear(in_features, out_features)
    def forward(self, x):
        logits = self.fc(x)
        return logits
    
classifier = Classifier(in_features = 5, out_features = 2)

bob.clear_objects()
classifier = classifier.fix_prec().share(bob, alice, crypto_provider = crypto_provider)
for i in range(3):
    a = torch.ones(1,5)
    b = a.fix_prec().share(bob, alice, crypto_provider = crypto_provider)
    classifier(b)
    print(len(bob._objects))
```

**Expected behavior**
the number of object in bob's store shouldn't increase, but currently it does (quite a lot)



Discussion:
I'd like to take up this issue!
And very importantly to note: if you comment out `classifier(b)` the problem is no longer there
Thanks @Syzygianinfern0!
Let us know if you're blocked. I've put the issue under "High Priority"
Suggestion: you might want to try first if simple multiplication has a GC issue :) 
> Thanks @Syzygianinfern0!
> Let us know if you're blocked. I've put the issue under "High Priority"

Sure. I'll ask for help if stuck.
> Suggestion: you might want to try first if simple multiplication has a GC issue :)

Looks like the normal arithemetic operations don't do that. Just the ones with the more complicated operations do it. Maybe the intermediate tensors during the computataion (the ones preserved for the backprop) are the ones with the problem. 

Resolving patch:
diff --git a/syft/frameworks/torch/tensors/interpreters/additive_shared.py b/syft/frameworks/torch/tensors/interpreters/additive_shared.py
--- a/syft/frameworks/torch/tensors/interpreters/additive_shared.py
+++ b/syft/frameworks/torch/tensors/interpreters/additive_shared.py
@@ -155,6 +155,7 @@ def get(self):
                 shares.append(share.get())
             else:
                 shares.append(share)
+                self.owner.de_register_obj(share)
 
         res_field = sum(shares) % self.field
 
@@ -247,7 +248,7 @@ def reconstruct(self):
         """
         workers = self.locations
 
-        ptr_to_sh = self.wrap().send(workers[0], **no_wrap)
+        ptr_to_sh = self.copy().wrap().send(workers[0], **no_wrap)
         pointer = ptr_to_sh.remote_get()
 
         pointers = [pointer]
@@ -1033,6 +1034,7 @@ def simplify(worker: AbstractWorker, tensor: "AdditiveSharingTensor") -> tuple:
             chain = sy.serde.msgpack.serde._simplify(worker, tensor.child)
 
         # Don't delete the remote values of the shares at simplification
+        garbage_collect = tensor.get_garbage_collect_data()
         tensor.set_garbage_collect_data(False)
 
         return (
@@ -1040,6 +1042,7 @@ def simplify(worker: AbstractWorker, tensor: "AdditiveSharingTensor") -> tuple:
             tensor.field,
             sy.serde.msgpack.serde._simplify(worker, tensor.crypto_provider.id),
             chain,
+            garbage_collect,
         )
 
     @staticmethod
@@ -1055,7 +1058,7 @@ def detail(worker: AbstractWorker, tensor_tuple: tuple) -> "AdditiveSharingTenso
                 shared_tensor = detail(data)
             """
 
-        tensor_id, field, crypto_provider, chain = tensor_tuple
+        tensor_id, field, crypto_provider, chain, garbage_collect = tensor_tuple
         crypto_provider = sy.serde.msgpack.serde._detail(worker, crypto_provider)
 
         tensor = AdditiveSharingTensor(
@@ -1069,6 +1072,8 @@ def detail(worker: AbstractWorker, tensor_tuple: tuple) -> "AdditiveSharingTenso
             chain = sy.serde.msgpack.serde._detail(worker, chain)
             tensor.child = chain
 
+        tensor.set_garbage_collect_data(garbage_collect)
+
         return tensor
 
     @staticmethod
diff --git a/syft/workers/base.py b/syft/workers/base.py
--- a/syft/workers/base.py
+++ b/syft/workers/base.py
@@ -542,13 +542,12 @@ def execute_communication_action(self, action: CommunicationAction) -> PointerTe
         else:
             obj = self.get_obj(obj_id)
             response = source_worker.send(obj, *destinations, **kwargs_)
-
-            response = hook_args.register_response("send", response, [sy.ID_PROVIDER.pop()], self)
-
-            # @lariffle: We only remove remote objects when the operations are inplace
-            # otherwise we could have stale pointers which we really want to avoid.
-            # TODO: needs more discussion
-            if kwargs_.get("inplace"):
+            if kwargs_.get("requires_grad", False):
+                response = hook_args.register_response(
+                    "send", response, [sy.ID_PROVIDER.pop()], self
+                )
+            else:
+                response.garbage_collect_data = False
                 self.rm_obj(obj_id)
             return response
 


Test patch:
diff --git a/test/serde/serde_helpers.py b/test/serde/serde_helpers.py
--- a/test/serde/serde_helpers.py
+++ b/test/serde/serde_helpers.py
@@ -534,6 +534,7 @@ def compare(detailed, original):
                     msgpack.serde._simplify(
                         kwargs["workers"]["serde_worker"], ast.child
                     ),  # (dict of AbstractTensor) simplified chain
+                    ast.get_garbage_collect_data(),
                 ),
             ),
             "cmp_detailed": compare,
diff --git a/test/torch/nn/test_nn.py b/test/torch/nn/test_nn.py
--- a/test/torch/nn/test_nn.py
+++ b/test/torch/nn/test_nn.py
@@ -7,6 +7,23 @@
 import syft.frameworks.torch.nn as syft_nn
 
 
+def test_nn_linear(workers):
+    torch.manual_seed(121)  # Truncation might not always work so we set the random seed
+    bob, alice, james = (workers["bob"], workers["alice"], workers["james"])
+    t = torch.tensor([[1.0, 2]])
+    x = t.fix_prec().share(bob, alice, crypto_provider=james)
+    model = nn.Linear(2, 1)
+    model.weight = nn.Parameter(torch.tensor([[-1.0, 2]]))
+    model.bias = nn.Parameter(torch.tensor([[-1.0]]))
+    model.fix_precision().share(bob, alice, crypto_provider=james)
+
+    y = model(x)
+
+    assert len(alice._objects) == 4  # x, y, weight, bias
+    assert len(bob._objects) == 4
+    assert y.get().float_prec() == torch.tensor([[2.0]])
+
+
 def test_conv2d(workers):
     """
     Test the nn.Conv2d module to ensure that it produces the exact same
diff --git a/test/torch/pointers/test_pointer_tensor.py b/test/torch/pointers/test_pointer_tensor.py
--- a/test/torch/pointers/test_pointer_tensor.py
+++ b/test/torch/pointers/test_pointer_tensor.py
@@ -289,7 +289,7 @@ def test_move(workers):
 
     x.move(alice)
 
-    assert x.id_at_location in bob._objects
+    assert x.id_at_location not in bob._objects
     assert x.id_at_location in alice._objects
 
     x = torch.tensor([1.0, 2, 3, 4, 5], requires_grad=True).send(bob)
@@ -299,7 +299,7 @@ def test_move(workers):
 
     x.move(alice)
 
-    assert x.id_at_location in bob._objects
+    assert x.id_at_location not in bob._objects
     assert x.id_at_location in alice._objects
 
     alice.clear_objects()
@@ -554,3 +554,15 @@ def test_setting_back_grad_to_origin_after_move(workers):
         z.backward()
 
         assert (x.grad == th.tensor([4.0, 4.0, 4.0, 4.0, 4.0])).all()
+
+
+def test_iadd(workers):
+    alice = workers["alice"]
+    a = torch.ones(1, 5)
+    b = torch.ones(1, 5)
+    a_pt = a.send(alice)
+    b_pt = b.send(alice)
+
+    b_pt += a_pt
+
+    assert len(alice._objects) == 2
diff --git a/test/torch/tensors/test_additive_shared.py b/test/torch/tensors/test_additive_shared.py
--- a/test/torch/tensors/test_additive_shared.py
+++ b/test/torch/tensors/test_additive_shared.py
@@ -493,21 +493,6 @@ def test_roll(workers):
     assert (res2.get() == torch.roll(t, (1, 2), dims=(0, 1))).all()
 
 
-def test_nn_linear(workers):
-    torch.manual_seed(121)  # Truncation might not always work so we set the random seed
-    bob, alice, james = (workers["bob"], workers["alice"], workers["james"])
-    t = torch.tensor([[1.0, 2]])
-    x = t.fix_prec().share(bob, alice, crypto_provider=james)
-    model = nn.Linear(2, 1)
-    model.weight = nn.Parameter(torch.tensor([[-1.0, 2]]))
-    model.bias = nn.Parameter(torch.tensor([[-1.0]]))
-    model.fix_precision().share(bob, alice, crypto_provider=james)
-
-    y = model(x)
-
-    assert y.get().float_prec() == torch.tensor([[2.0]])
-
-
 def test_matmul(workers):
     torch.manual_seed(121)  # Truncation might not always work so we set the random seed
     bob, alice, james = (workers["bob"], workers["alice"], workers["james"])
@@ -859,3 +844,37 @@ def test_correct_tag_and_description_after_send(workers):
 
     assert me.request_search("tag_additive_test1", location=alice)
     assert me.request_search("tag_additive_test2", location=alice)
+
+
+def test_garbage_collect_reconstruct(workers):
+    bob, alice, james, me = (workers["bob"], workers["alice"], workers["james"], workers["me"])
+    a = torch.ones(1, 5)
+    a_sh = a.encrypt(workers=[alice, bob], crypto_provider=james)
+    a_recon = a_sh.child.child.reconstruct()
+
+    assert len(alice._objects) == 2
+    assert len(bob._objects) == 2
+
+
+def test_garbage_collect_move(workers):
+    bob, alice, me = (workers["bob"], workers["alice"], workers["me"])
+    a = torch.ones(1, 5).send(alice)
+    b = a.copy().move(bob)
+
+    assert len(alice._objects) == 1
+    assert len(bob._objects) == 1
+
+
+def test_garbage_collect_mul(workers):
+    bob, alice, james, me = (workers["bob"], workers["alice"], workers["james"], workers["me"])
+    a = torch.ones(1, 5)
+    b = torch.ones(1, 5)
+
+    a = a.encrypt(workers=[alice, bob], crypto_provider=james)
+    b = b.encrypt(workers=[alice, bob], crypto_provider=james)
+
+    for _ in range(3):
+        c = a * b
+
+    assert len(alice._objects) == 3
+    assert len(bob._objects) == 3