Repository: OpenMined/PySyft
Pull request: Fix Memory Leak in AST Operations
Fixes #3316 

## Current progress
- [x] Worker 1
- [x] Worker 2
- [x] Crypto Provider
- [x] Multi Party
- [x] Clean-Up
- [x] Tests


## The Change in Behaviour

<details>
  <summary>Code Used to Test</summary>
  
```python
import torch as th
import syft as sy

hook = sy.TorchHook(th)
alice = sy.VirtualWorker(hook, id="alice")
bob = sy.VirtualWorker(hook, id="bob")
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


def mul(a, b):
    c = a * b
    return c


def add(a, b):
    c = a + b
    return c


alice.clear_objects()
bob.clear_objects()

a = torch.ones(1, 5)
b = torch.ones(1, 5)

a = a.encrypt(workers=[alice, bob], crypto_provider=crypto_provider)
b = b.encrypt(workers=[alice, bob], crypto_provider=crypto_provider)


classifier = Classifier(in_features=5, out_features=2)
classifier = classifier.fix_prec().share(bob, alice, crypto_provider=crypto_provider)


for i in range(3):
    print("-" * 20)
    print(f"Alice: {len(alice._objects)}")
    print(f"Bob: {len(bob._objects)}")
    print(f"Crypto Provider: {len(crypto_provider._objects)}")

    c = classifier(b)  # <<<<< Just Toggle Between Operations
    # c = mul(a, b)
    # c = add(a, b)

    print(f"Alice: {len(alice._objects)}")
    print(f"Bob: {len(bob._objects)}")
    print(f"crypto_provider: {len(crypto_provider._objects)}")

```
</details>

![Test](https://i.imgur.com/noAPMr5.png)

<details>
<summary>Same Table in MD</summary>

#### Sorry for the Ugly Table 😅
(\<br>'s don't go well in GitHub)
<table>
    <tr>
        <th>Operation</th>
        <th>Current Behaviour</th>
        <th>After Merge</th>
    </tr>
    <tr>
        <td>Addition</td>
        <td>
            <code>==========Iter1==========<br>
            Alice: 4<br>
            Bob: 4<br>
            Crypto Provider: 0<br>
            !!OPERATION!!<br>
            Alice: 5<br>
            Bob: 5<br>
            crypto_provider: 0<br>
            ==========Iter2==========<br>
            Alice: 5<br>
            Bob: 5<br>
            Crypto Provider: 0<br>
            !!OPERATION!!<br>
            Alice: 5<br>
            Bob: 5<br>
            crypto_provider: 0<br>
            ==========Iter3==========<br>
            Alice: 5<br>
            Bob: 5<br>
            Crypto Provider: 0<br>
            !!OPERATION!!<br>
            Alice: 5<br>
            Bob: 5<br>
            crypto_provider: 0</code>
        </td>
        <td>
            <code>==========Iter1==========<br>
            Alice: 4<br>
            Bob: 4<br>
            Crypto Provider: 0<br>
            !!OPERATION!!<br>
            Alice: 5<br>
            Bob: 5<br>
            crypto_provider: 0<br>
            ==========Iter2==========<br>
            Alice: 5<br>
            Bob: 5<br>
            Crypto Provider: 0<br>
            !!OPERATION!!<br>
            Alice: 5<br>
            Bob: 5<br>
            crypto_provider: 0<br>
            ==========Iter3==========<br>
            Alice: 5<br>
            Bob: 5<br>
            Crypto Provider: 0<br>
            !!OPERATION!!<br>
            Alice: 5<br>
            Bob: 5<br>
            crypto_provider: 0
            </code>
        </td>
    </tr>
        <tr>
        <td>Multiplication</td>
        <td>
            <code>==========Iter1==========<br>
                Alice: 4<br>
                Bob: 4<br>
                Crypto Provider: 0<br>
                !!OPERATION!!<br>
                Alice: 12<br>
                Bob: 6<br>
                crypto_provider: 1<br>
                ==========Iter2==========<br>
                Alice: 12<br>
                Bob: 6<br>
                Crypto Provider: 1<br>
                !!OPERATION!!<br>
                Alice: 19<br>
                Bob: 7<br>
                crypto_provider: 2<br>
                ==========Iter3==========<br>
                Alice: 19<br>
                Bob: 7<br>
                Crypto Provider: 2<br>
                !!OPERATION!!<br>
                Alice: 26<br>
                Bob: 8<br>
                crypto_provider: 3            
            </code>
        </td>
        <td>
            <code>==========Iter1==========<br>
            Alice: 4<br>
            Bob: 4<br>
            Crypto Provider: 0<br>
            !!OPERATION!!<br>
            Alice: 5<br>
            Bob: 5<br>
            crypto_provider: 0<br>
            ==========Iter2==========<br>
            Alice: 5<br>
            Bob: 5<br>
            Crypto Provider: 0<br>
            !!OPERATION!!<br>
            Alice: 5<br>
            Bob: 5<br>
            crypto_provider: 0<br>
            ==========Iter3==========<br>
            Alice: 5<br>
            Bob: 5<br>
            Crypto Provider: 0<br>
            !!OPERATION!!<br>
            Alice: 5<br>
            Bob: 5<br>
            crypto_provider: 0
            </code>
        </td>
    </tr>
    <tr>
        <td>nn.Linear</td>
        <td>
            <code>==========Iter1==========<br>
            Alice: 4<br>
            Bob: 4<br>
            Crypto Provider: 0<br>
            !!OPERATION!!<br>
            Alice: 9<br>
            Bob: 9<br>
            crypto_provider: 1<br>
            ==========Iter2==========<br>
            Alice: 9<br>
            Bob: 9<br>
            Crypto Provider: 1<br>
            !!OPERATION!!<br>
            Alice: 13<br>
            Bob: 13<br>
            crypto_provider: 2<br>
            ==========Iter3==========<br>
            Alice: 13<br>
            Bob: 13<br>
            Crypto Provider: 2<br>
            !!OPERATION!!<br>
            Alice: 17<br>
            Bob: 17<br>
            crypto_provider: 3
            </code>
        </td>
        <td>
            <code>
            ==========Iter1==========<br>
            Alice: 4<br>
            Bob: 4<br>
            Crypto Provider: 0<br>
            !!OPERATION!!<br>
            Alice: 5<br>
            Bob: 5<br>
            crypto_provider: 0<br>
            ==========Iter2==========<br>
            Alice: 5<br>
            Bob: 5<br>
            Crypto Provider: 0<br>
            !!OPERATION!!<br>
            Alice: 5<br>
            Bob: 5<br>
            crypto_provider: 0<br>
            ==========Iter3==========<br>
            Alice: 5<br>
            Bob: 5<br>
            Crypto Provider: 0<br>
            !!OPERATION!!<br>
            Alice: 5<br>
            Bob: 5<br>
            crypto_provider: 0
            </code>
        </td>
    </tr>
</table>
</details>

Problem statement:
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)

Hints from issue 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. 
It seems there is an issue with the mul operation in AdditiveSharedTensor. It maintains copies of the tensors on which the operation is performed.

Commit message: Add move, mul and nn tests

Patch to review:
diff --git a/syft/frameworks/torch/tensors/interpreters/additive_shared.py b/syft/frameworks/torch/tensors/interpreters/additive_shared.py
index 5533508170c..da0dcafb7d8 100644
--- 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/generic/frameworks/hook/hook_args.py b/syft/generic/frameworks/hook/hook_args.py
index 3d3957ff1ac..c7fbc309b76 100644
--- a/syft/generic/frameworks/hook/hook_args.py
+++ b/syft/generic/frameworks/hook/hook_args.py
@@ -711,6 +711,7 @@ def register_tensor(
         response_ids: List of ids where the tensor should be stored
             and each id is pop out when needed.
     """
+    # owner.de_register_obj(tensor)
     tensor.owner = owner
     try:
         tensor.id = response_ids.pop(-1)
diff --git a/syft/workers/base.py b/syft/workers/base.py
index 467e88e77db..76f7384dab0 100644
--- 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
 
diff --git a/test/serde/serde_helpers.py b/test/serde/serde_helpers.py
index 6e011d0c1d4..cf378039b0a 100644
--- a/test/serde/serde_helpers.py
+++ b/test/serde/serde_helpers.py
@@ -532,6 +532,7 @@ def compare(detailed, original):
                     msgpack.serde._simplify(
                         syft.hook.local_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
index e67506a582d..9c6baa4c90e 100644
--- a/test/torch/nn/test_nn.py
+++ b/test/torch/nn/test_nn.py
@@ -8,6 +8,23 @@
 import syft.frameworks.torch.nn as nn2
 
 
+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
index f71e5cfe370..a66d8620f58 100644
--- 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 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 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
index 0f0c1ace997..76c4f0727eb 100644
--- 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


Code review comments:
diff_hunk: @@ -711,6 +711,7 @@ def register_tensor(
         response_ids: List of ids where the tensor should be stored
             and each id is pop out when needed.
     """
+    # owner.de_register_obj(tensor)
path: syft/generic/frameworks/hook/hook_args.py
text: @user1:
That's not needed finally?

@author:
Yes. Works without it.
Removed it as it was breaking the inplace operations such as ```__iadd__```, etc

@user2:
If it's not needed, let's remove it.

@author:
@user2 Done