Repository: lightly-ai/lightly

Issue:
NNMemoryBank not working with DataParallel
I have been using the NNMemoryBank as a component in my module and noticed that at each forward pass `NNMemoryBank.bank` is equal to `None`. This only occurs when my module is wrapped in `DataParallel`. As a result, throughout training my NN pairs are always random noise (surprisingly, this only hurt the contrastive learning performance by a few percentage point on linear probing??).

Here is a simple test case that highlights the issue:
```
import torch
from lightly.models.modules import NNMemoryBankModule
memory_bank = NNMemoryBankModule(size=1000)
print(memory_bank.bank)
memory_bank(torch.randn((100, 512)))
print(memory_bank.bank)

memory_bank = NNMemoryBankModule(size=1000)
memory_bank = torch.nn.DataParallel(memory_bank, device_ids=[0,1])
print(memory_bank.module.bank)
memory_bank(torch.randn((100, 512)))
print(memory_bank.module.bank)
```

The output of the first is `None` and a `torch.Tensor`, as expected. The output for the second is `None` for both.


Discussion:
Hi @kfallah, thanks for reporting this! I'll have to take a closer look to figure out what's wrong there but let me get back to you 🙂 In the meantime: If you have any idea what might be happening please just drop a comment below.
I just wanted to check back and see if there was any progress on this issue 🙂. Could this be fixed by making `bank` a parameter with `requires_grad=False` in `lightly.loss.memory_bank.MemoryBankModule`? I think this same issue may also occur with loading/saving the module, since `bank` is just a tensor.
Note that this issue is a known limitation in the DataParallel, where the suggested solution is an in-place edit of a buffer or parameter (second warning blurb as of Nov 2022):
https://pytorch.org/docs/stable/generated/torch.nn.DataParallel.html 

Seems like the best design pattern would be to register `bank` as a buffer and do in-place edits in `forward`:
https://pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module.register_buffer
https://discuss.pytorch.org/t/what-is-the-difference-between-register-buffer-and-register-parameter-of-nn-module/32723

But also seems like there was a deliberate decision not to do this:
https://github.com/lightly-ai/lightly/blob/master/lightly/loss/memory_bank.py#L62
Would an extra `66536 x 256 = ~16 MB` really make a big difference (`65536` is the memory bank size from NNCLR, `256` for the projected feature dimension)?
Hi @kfallah,

Thanks a lot for investigating this! I think you are right and we should register the memory bank as a buffer. `register_buffer` takes an additional `persistent` argument which controls whether the buffer should be part of the state dict or not. I believe we can just set `persistent=False` to avoid saving the memory bank in the checkpoint.

Would you be interested in making a PR for this?
I do not think I will have the time to make the PR, I would be deeply grateful if someone else had the chance.
I'll have a look as soon as I get some time.

Resolving patch:
diff --git a/lightly/loss/memory_bank.py b/lightly/loss/memory_bank.py
--- a/lightly/loss/memory_bank.py
+++ b/lightly/loss/memory_bank.py
@@ -46,10 +46,9 @@ def __init__(self, size: int = 2 ** 16):
             raise ValueError(msg)
 
         self.size = size
+        self.register_buffer("bank", tensor=torch.empty(0, dtype=torch.float), persistent=False)
+        self.register_buffer("bank_ptr", tensor=torch.empty(0, dtype=torch.long), persistent=False)
 
-        self.bank = None
-        self.bank_ptr = None
-    
     @torch.no_grad()
     def _init_memory_bank(self, dim: int):
         """Initialize the memory bank if it's empty
@@ -63,9 +62,9 @@ def _init_memory_bank(self, dim: int):
         # we could use register buffers like in the moco repo
         # https://github.com/facebookresearch/moco but we don't
         # want to pollute our checkpoints
-        self.bank = torch.randn(dim, self.size)
-        self.bank = torch.nn.functional.normalize(self.bank, dim=0)
-        self.bank_ptr = torch.LongTensor([0])
+        self.bank = torch.randn(dim, self.size).type_as(self.bank)
+        torch.nn.functional.normalize(self.bank, dim=0)
+        self.bank_ptr = torch.zeros(1).type_as(self.bank_ptr)
 
     @torch.no_grad()
     def _dequeue_and_enqueue(self, batch: torch.Tensor):
@@ -111,7 +110,7 @@ def forward(self,
         _, dim = output.shape
 
         # initialize the memory bank if it is not already done
-        if self.bank is None:
+        if self.bank.nelement() == 0:
             self._init_memory_bank(dim)
 
         # query and update memory bank


Test patch:
diff --git a/tests/loss/test_MemoryBank.py b/tests/loss/test_MemoryBank.py
--- a/tests/loss/test_MemoryBank.py
+++ b/tests/loss/test_MemoryBank.py
@@ -6,11 +6,11 @@
 
 class TestNTXentLoss(unittest.TestCase):
 
-    def test_NegativeSize(self):
+    def test_init__negative_size(self):
         with self.assertRaises(ValueError):
             MemoryBankModule(size=-1)
 
-    def test_ForwardEasy(self):
+    def test_forward_easy(self):
         bsz = 3
         dim, size = 2, 9
         n = 33 * bsz
@@ -36,7 +36,7 @@ def test_ForwardEasy(self):
 
             ptr = (ptr + bsz) % size
 
-    def test_Forward(self):
+    def test_forward(self):
         bsz = 3
         dim, size = 2, 10
         n = 33 * bsz
@@ -48,3 +48,19 @@ def test_Forward(self):
             # is no multiple of the batch size
             output = torch.randn(bsz, dim)
             _, _ = memory_bank(output)
+
+    @unittest.skipUnless(torch.cuda.is_available(), "cuda not available")
+    def test_forward__cuda(self):
+        bsz = 3
+        dim, size = 2, 10
+        n = 33 * bsz
+        memory_bank = MemoryBankModule(size=size)
+        device = torch.device('cuda')
+        memory_bank.to(device=device)
+
+        for i in range(0, n, bsz):
+
+            # see if there are any problems when the bank size
+            # is no multiple of the batch size
+            output = torch.randn(bsz, dim, device=device)
+            _, _ = memory_bank(output)