Pysyft: Garbage collection: memory leak in SMPC

[LaRiffle]

Describe the bug
Garbage collection is not working as expected in SMPC

To Reproduce

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)

[Syzygianinfern0]

I'd like to take up this issue!

[AlanAboudib]

And very importantly to note: if you comment out classifier(b) the problem is no longer there

[LaRiffle]

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

[Syzygianinfern0]

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.

[Syzygianinfern0]

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.

[sachin-101]

It seems there is an issue with the mul operation in AdditiveSharedTensor. It maintains copies of the tensors on which the operation is performed.