Repository: ElementAI/baal
Pull request #7 (2019-10-11 17:11:43)

Cuda mem overflow
## Summary:
When running the predict_* methods of ModelWrapper, the same batch size that works during training / testing may cause out-of-memory errors because we replicate the data in each batch so the memory footprint will be e.g. 10x the memory footprint during training.

### Features:
Added a replicate_in_memory instantiation parameter which allows for out-of-memory data batch replication.

## Checklist:

* [*] Your code is documented (To validate this, add your module to `tests/documentation_test.py`).
* [*] Your code is tested with unit tests.
* [*] You moved your Issue to the PR state.



Patch:
diff --git a/src/baal/modelwrapper.py b/src/baal/modelwrapper.py
index ad2128f..87aecb6 100644
--- a/src/baal/modelwrapper.py
+++ b/src/baal/modelwrapper.py
@@ -25,13 +25,15 @@ class ModelWrapper:
     Args:
         model (nn.Module): The model to optimize.
         criterion (Callable): a loss function.
+        replicate_in_memory (bool): Replicate in memory optional.
     """
 
-    def __init__(self, model, criterion):
+    def __init__(self, model, criterion, replicate_in_memory=True):
         self.model = model
         self.criterion = criterion
         self.metrics = dict()
         self.add_metric('loss', lambda: Loss())
+        self.replicate_in_memory = replicate_in_memory
 
     def add_metric(self, name: str, initializer: Callable):
         """
@@ -335,17 +337,39 @@ class ModelWrapper:
         Returns:
             Tensor, the loss computed from the criterion.
                     shape = {batch_size, nclass, n_iteration}
+
+        Raises:
+            raises RuntimeError if CUDA rans out of memory during data replication.
         """
         with torch.no_grad():
             if cuda:
                 data = data.cuda()
-            input_shape = data.size()
-            batch_size = input_shape[0]
-            data = torch.stack([data] * iterations)
-            data = data.view(batch_size * iterations, *input_shape[1:])
-            out = self.model(data)
-            out = map_on_tensor(lambda o: o.view([iterations, batch_size, *o.size()[1:]]), out)
-            out = map_on_tensor(lambda o: o.permute(1, 2, *range(3, o.ndimension()), 0), out)
+            if self.replicate_in_memory:
+                input_shape = data.size()
+                batch_size = input_shape[0]
+                try:
+                    data = torch.stack([data] * iterations)
+                except RuntimeError as e:
+                    raise RuntimeError(
+                        '''CUDA ran out of memory while BaaL tried to replicate data. See the exception above.
+                    Use `replicate_in_memory=False` in order to reduce the memory requirements.
+                    Note that there will be some speed trade-offs''') from e
+                data = data.view(batch_size * iterations, *input_shape[1:])
+                try:
+                    out = self.model(data)
+                except RuntimeError as e:
+                    raise RuntimeError(
+                        '''CUDA ran out of memory while BaaL tried to replicate data. See the exception above.
+                    Use `replicate_in_memory=False` in order to reduce the memory requirements.
+                    Note that there will be some speed trade-offs''') from e
+                out = map_on_tensor(lambda o: o.view([iterations, batch_size, *o.size()[1:]]), out)
+                out = map_on_tensor(lambda o: o.permute(1, 2, *range(3, o.ndimension()), 0), out)
+            else:
+                out = [self.model(data) for _ in range(iterations)]
+                if isinstance(out[0], Sequence):
+                    out = [torch.stack(ts, dim=-1) for ts in zip(*out)]
+                else:
+                    out = torch.stack(out, dim=-1)
             return out
 
     def get_params(self):


Test patch:
diff --git a/tests/modelwrapper_test.py b/tests/modelwrapper_test.py
index 702d3c2..a142900 100644
--- a/tests/modelwrapper_test.py
+++ b/tests/modelwrapper_test.py
@@ -1,3 +1,4 @@
+import math
 import unittest
 from unittest.mock import Mock
 
@@ -99,6 +100,36 @@ class ModelWrapperMultiOutTest(unittest.TestCase):
         pred = self.wrapper.predict_on_batch(input, 10, False)
         assert pred[0].size() == (2, 1, 10)
 
+        # iteration == 1
+        self.wrapper = ModelWrapper(self.model, self.criterion, replicate_in_memory=False)
+        pred = self.wrapper.predict_on_batch(input, 1, False)
+        assert pred[0].size() == (2, 1, 1)
+
+        # iterations > 1
+        pred = self.wrapper.predict_on_batch(input, 10, False)
+        assert pred[0].size() == (2, 1, 10)
+
+    def test_out_of_mem_raises_error(self):
+        self.wrapper.eval()
+        input = torch.stack((self.dataset[0][0], self.dataset[1][0]))
+        with pytest.raises(RuntimeError) as e_info:
+            self.wrapper.predict_on_batch(input, 0, False)
+        assert 'CUDA ran out of memory while BaaL tried to replicate data' in str(e_info.value)
+
+    def test_raising_type_errors(self):
+        iterations = math.inf
+        self.wrapper.eval()
+        input = torch.stack((self.dataset[0][0], self.dataset[1][0]))
+        with pytest.raises(TypeError):
+            self.wrapper.predict_on_batch(input, iterations, False)
+
+    def test_using_cuda_raises_error_while_testing(self):
+        '''CUDA is not available on test environment'''
+        self.wrapper.eval()
+        input = torch.stack((self.dataset[0][0], self.dataset[1][0]))
+        with pytest.raises(Exception):
+            self.wrapper.predict_on_batch(input, 1, True)
+
     def test_train(self):
         history = self.wrapper.train_on_dataset(self.dataset, self.optim, 10, 2, use_cuda=False,
                                                 workers=0)
@@ -140,6 +171,12 @@ class ModelWrapperMultiOutTest(unittest.TestCase):
         assert next(l_gen)[0].dtype == np.float16
         assert l[0].dtype == np.float16
 
+        data_s = []
+        l_gen = self.wrapper.predict_on_dataset_generator(data_s, 10, 20, use_cuda=False,
+                                                          workers=0, half=True)
+
+        assert len(list(l_gen)) == 0
+
 
 class ModelWrapperTest(unittest.TestCase):
     def setUp(self):
@@ -207,6 +244,15 @@ class ModelWrapperTest(unittest.TestCase):
         pred = self.wrapper.predict_on_batch(input, 10, False)
         assert pred.size() == (2, 1, 10)
 
+        # iteration == 1
+        self.wrapper = ModelWrapper(self.model, self.criterion, replicate_in_memory=False)
+        pred = self.wrapper.predict_on_batch(input, 1, False)
+        assert pred.size() == (2, 1, 1)
+
+        # iterations > 1
+        pred = self.wrapper.predict_on_batch(input, 10, False)
+        assert pred.size() == (2, 1, 10)
+
     def test_train(self):
         history = self.wrapper.train_on_dataset(self.dataset, self.optim, 10, 2, use_cuda=False,
                                                 workers=0)
@@ -248,26 +294,39 @@ class ModelWrapperTest(unittest.TestCase):
         assert l.dtype == np.float16
 
     def test_states(self):
-        self.wrapper.train()
         input = torch.randn([1, 3, 10, 10])
-        # Dropout make the pred changes
-        preds = torch.stack(
-            [
-                self.wrapper.predict_on_batch(input, iterations=1, cuda=False)
-                for _ in range(10)
-            ]
-        ).view(10, -1)
-        assert not torch.allclose(torch.mean(preds, 0), preds[0])
 
-        # Dropout is not active in eval
-        self.wrapper.eval()
-        preds = torch.stack(
-            [
-                self.wrapper.predict_on_batch(input, iterations=1, cuda=False)
-                for _ in range(10)
-            ]
-        ).view(10, -1)
-        assert torch.allclose(torch.mean(preds, 0), preds[0])
+        def pred_with_dropout(replicate_in_memory):
+            self.wrapper = ModelWrapper(self.model, self.criterion,
+                                        replicate_in_memory=replicate_in_memory)
+            self.wrapper.train()
+            # Dropout make the pred changes
+            preds = torch.stack(
+                [
+                    self.wrapper.predict_on_batch(input, iterations=1, cuda=False)
+                    for _ in range(10)
+                ]
+            ).view(10, -1)
+            assert not torch.allclose(torch.mean(preds, 0), preds[0])
+
+        pred_with_dropout(replicate_in_memory=True)
+        pred_with_dropout(replicate_in_memory=False)
+
+        def pred_without_dropout(replicate_in_memory):
+            self.wrapper = ModelWrapper(self.model, self.criterion,
+                                        replicate_in_memory=replicate_in_memory)
+            # Dropout is not active in eval
+            self.wrapper.eval()
+            preds = torch.stack(
+                [
+                    self.wrapper.predict_on_batch(input, iterations=1, cuda=False)
+                    for _ in range(10)
+                ]
+            ).view(10, -1)
+            assert torch.allclose(torch.mean(preds, 0), preds[0])
+
+        pred_without_dropout(replicate_in_memory=True)
+        pred_without_dropout(replicate_in_memory=False)
 
     def test_add_metric(self):
         self.wrapper.add_metric('cls_report', lambda: ClassificationReport(2))
@@ -309,6 +368,5 @@ class ModelWrapperTest(unittest.TestCase):
         assert len(res[0]) < 50 and len(res[0]) > 20
 
 
-
 if __name__ == '__main__':
     pytest.main()