Repository: OpenMined/PySyft
Pull request: Garbage Collection in `execute_computation_action`, Fix Inplace-Operation Identification
## Description

Fixes #3397 
Fixes #3398 

~Currently breaks inplace computations. WIP 🚧🚧~
- [x] GC
- [x] Fix inplace operations

## Type of change

Please mark options that are relevant.

- [ ] Added/Modified tutorials
- [x] Bug fix (non-breaking change which fixes an issue)
- [ ] New feature (non-breaking change which adds functionality)
- [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected)

## Checklist:

* [x] I have added tests for my changes
* [x] I did follow the [contribution guidelines](https://github.com/OpenMined/PySyft/blob/master/CONTRIBUTING.md)
* [x] I have commented my code following [Google style](https://sphinxcontrib-napoleon.readthedocs.io/en/latest/example_google.html).

Problem statement:
Memory leak in F.relu() method in AST 
This code, causes a memory leak only on the first worker in the worker list passed to the `encrypt()`. In this example it is  'alice', the second worker 'bob' does not suffer from any leaks.

The reason is the `F.relu()` method:

```
import torch as th
import torch.optim as optim
import pprint
import syft as sy
import torch.nn.functional as F

hook = sy.TorchHook(th) #, verbose=True)
bob = sy.VirtualWorker(hook, id="bob") #, verbose=True)
alice = sy.VirtualWorker(hook, id="alice") #, verbose=True)
crypto_provider = sy.VirtualWorker(hook, id="james") #, verbose=True)



a = th.ones(1, 5)
a = a.encrypt(workers=[alice, bob], crypto_provider=crypto_provider, requires_grad=True)


class Classifier(th.nn.Module):
    def __init__(self, in_features, out_features):
        super(Classifier, self).__init__()
        self.fc = th.nn.Linear(in_features, out_features)

    def forward(self, x):
        logits = self.fc(x)
        logits = F.relu(logits)
        return logits


classifier = Classifier(in_features=5, out_features=5)
model = classifier.fix_prec().share(alice, bob, crypto_provider=crypto_provider, requires_grad=True)

for i in range(3):
    print("=" * 10 + f"Iter{i + 1}" + "=" * 10)
    print(f"Alice: {len(alice._objects)}")
    print(f"Bob: {len(bob._objects)}")

    preds = classifier(a)


    print(f"Alice: {len(alice._objects)}")
    print(f"Bob: {len(bob._objects)}")

print("~" * 10 + "Done" + "~" * 10)
```

Here is the output (Only Alice leaks)

```
==========Iter1==========
Alice: 3
Bob: 3
Alice: 11
Bob: 7
==========Iter2==========
Alice: 11
Bob: 7
Alice: 15
Bob: 7
==========Iter3==========
Alice: 15
Bob: 7
Alice: 19
Bob: 7
~~~~~~~~~~Done~~~~~~~~~~

```


If you comment out `logits = F.relu(logits)`, no memory leaks are observed:

```
==========Iter1==========
Alice: 3
Bob: 3
Alice: 3
Bob: 3
==========Iter2==========
Alice: 3
Bob: 3
Alice: 3
Bob: 3
==========Iter3==========
Alice: 3
Bob: 3
Alice: 3
Bob: 3
~~~~~~~~~~Done~~~~~~~~~~

```

Memory leak in the argmax() method of AST
This code, causes a memory leak only on the first worker in the worker list passed to the `encrypt()`. In this example it is  'alice', the second worker 'bob' does not suffer from any leaks.

The reason is the `argmax()` method:

```
import torch as th
import torch.optim as optim
import syft as sy

hook = sy.TorchHook(th) #, verbose=True)
bob = sy.VirtualWorker(hook, id="bob") #, verbose=True)
alice = sy.VirtualWorker(hook, id="alice") #, verbose=True)
crypto_provider = sy.VirtualWorker(hook, id="james") #, verbose=True)

a = th.ones(1, 5)
a = a.encrypt(workers=[alice, bob], crypto_provider=crypto_provider, requires_grad=True)

for i in range(3):
    print("=" * 10 + f"Iter{i + 1}" + "=" * 10)
    print(f"Alice: {len(alice._objects)}")
    print(f"Bob: {len(bob._objects)}")

    a.argmax(dim=1)

    print(f"Alice: {len(alice._objects)}")
    print(f"Bob: {len(bob._objects)}")

print("~" * 10 + "Done" + "~" * 10)
```

Here is the output (Only Alice leaks):


```
==========Iter1==========
Alice: 1
Bob: 1
Alice: 13
Bob: 1
==========Iter2==========
Alice: 13
Bob: 1
Alice: 25
Bob: 1
==========Iter3==========
Alice: 25
Bob: 1
Alice: 37
Bob: 1
~~~~~~~~~~Done~~~~~~~~~~
```

If you comment out `a.argmax(dim=1)`, no memory leaks are observed:

```
==========Iter1==========
Alice: 1
Bob: 1
Alice: 1
Bob: 1
==========Iter2==========
Alice: 1
Bob: 1
Alice: 1
Bob: 1
==========Iter3==========
Alice: 1
Bob: 1
Alice: 1
Bob: 1
~~~~~~~~~~Done~~~~~~~~~~



```

Hints from issue discussion:

Commit message: Cleanup

Patch to review:
diff --git a/syft/frameworks/torch/tensors/interpreters/additive_shared.py b/syft/frameworks/torch/tensors/interpreters/additive_shared.py
index 0fb821e30bd..1e0535ff71b 100644
--- a/syft/frameworks/torch/tensors/interpreters/additive_shared.py
+++ b/syft/frameworks/torch/tensors/interpreters/additive_shared.py
@@ -943,7 +943,7 @@ def max(self, dim=None, return_idx=False):
             the maximum value (possibly across an axis)
             and optionally the index of the maximum value (possibly across an axis)
         """
-        values = self
+        values = self.copy()
         n_dim = self.dim()
 
         # Make checks and transformation
diff --git a/syft/frameworks/torch/torch_attributes.py b/syft/frameworks/torch/torch_attributes.py
index 946f35b215e..f2cf777ba74 100644
--- a/syft/frameworks/torch/torch_attributes.py
+++ b/syft/frameworks/torch/torch_attributes.py
@@ -1,10 +1,8 @@
+import re
 from types import ModuleType
-from typing import Union
-from typing import Callable
-from typing import Any
 
-from syft.generic.frameworks.attributes import FrameworkAttributes
 from syft.frameworks.torch.tensors.interpreters.native import TorchTensor
+from syft.generic.frameworks.attributes import FrameworkAttributes
 
 
 class TorchAttributes(FrameworkAttributes):
@@ -135,6 +133,12 @@ def __init__(self, torch: ModuleType, hook: ModuleType) -> None:
 
         # Dict {method_name: <is_inplace:bool>
         self.inplace_methods = {}
+        self._inplace_pattern = re.compile(r"(^__i(?!nit|mport|ter).+_)|^((?!^_+).+)_$")
+        # Positives:
+        # __iadd__, share_
+
+        # Negatives:
+        # __init__, __import__, __iter__, __foo__, __bar_foo
 
     def is_inplace_method(self, method_name):
         """Determine if a method is inplace or not.
@@ -150,6 +154,7 @@ def is_inplace_method(self, method_name):
         try:
             return self.inplace_methods[method_name]
         except KeyError:
-            is_inplace = method_name[-1] == "_" and "__" not in method_name
+            is_inplace = True if re.search(self._inplace_pattern, method_name) else False
+
             self.inplace_methods[method_name] = is_inplace
             return is_inplace
diff --git a/syft/generic/frameworks/hook/hook_args.py b/syft/generic/frameworks/hook/hook_args.py
index 17ff5fcdb62..9850c3568ec 100644
--- a/syft/generic/frameworks/hook/hook_args.py
+++ b/syft/generic/frameworks/hook/hook_args.py
@@ -717,6 +717,9 @@ def register_tensor(
         response_ids: List of ids where the tensor should be stored
             and each id is pop out when needed.
     """
+    # `de_register` initially as this causes the original PT
+    # to not get GC'd. Inplace operations do not follow this as expected.
+    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 301e5301a47..a09f0755bed 100644
--- a/syft/workers/base.py
+++ b/syft/workers/base.py
@@ -492,6 +492,7 @@ def execute_computation_action(self, action: ComputationAction) -> PointerTensor
             if type(_self) == str and _self == "self":
                 _self = self
             if sy.framework.is_inplace_method(op_name):
+                # TODO [sharan]: Can this be used to fix the inplace operations breakages??
                 # TODO[jvmancuso]: figure out a good way to generalize the
                 # above check (#2530)
                 getattr(_self, op_name)(*args_, **kwargs_)


Code review comments:
diff_hunk: @@ -943,7 +943,7 @@ def max(self, dim=None, return_idx=False):
             the maximum value (possibly across an axis)
             and optionally the index of the maximum value (possibly across an axis)
         """
-        values = self
+        values = self.copy()
path: syft/frameworks/torch/tensors/interpreters/additive_shared.py
text: @user1:
Does this inherently require copying because operations later in this function are in-place? Could it be avoided by changing the lines below that read something like `values = values.some_op()`?

@author:
Removed (redundant operation)

diff_hunk: @@ -492,6 +492,7 @@ def execute_computation_action(self, action: ComputationAction) -> PointerTensor
             if type(_self) == str and _self == "self":
                 _self = self
             if sy.framework.is_inplace_method(op_name):
+                # TODO [sharan]: Can this be used to fix the inplace operations breakages??
path: syft/workers/base.py
text: @user1:
If there's work left to be done, better to create an issue than a TODO.

@author:
Removed (it was addressed)