feature/1124 Fixed memory consumption in PL tests

airctic · Jul 22, 2022 · 8216315 · 8216315
1 parent 8f39eff
commit 8216315
Show file tree

Hide file tree

Showing 19 changed files with 40 additions and 0 deletions.
diff --git a/tests/models/efficient_det/lightning/test_test.py b/tests/models/efficient_det/lightning/test_test.py
@@ -47,6 +47,7 @@ def test_lightining_efficientdet_finalizes_metrics_on_test_epoch_end(
     fridge_efficientdet_model, light_model_cls, metrics
 ):
     light_model = light_model_cls(fridge_efficientdet_model, metrics=metrics)
+    torch.set_grad_enabled(False)
 
     light_model.test_epoch_end(None)
 
@@ -63,6 +64,7 @@ def test_lightining_efficientdet_logs_losses_during_test_step(
     light_model.convert_raw_predictions = lambda *args: None
     light_model.compute_loss = lambda *args: None
     light_model.accumulate_metrics = lambda *args: None
+    torch.set_grad_enabled(False)
 
     light_model.test_step(batch, 0)
 

diff --git a/tests/models/efficient_det/lightning/test_train.py b/tests/models/efficient_det/lightning/test_train.py
@@ -53,6 +53,7 @@ def fake_compute_loss(self, *args):
         return expected_loss
 
     light_model.compute_loss = fake_compute_loss
+    torch.set_grad_enabled(False)
 
     loss = light_model.training_step(batch, 0)
 
@@ -66,6 +67,7 @@ def test_lightining_efficientdet_logs_losses_during_training_step(
     light_model = light_model_cls(model=fridge_efficientdet_model, metrics=None)
     for batch in train_dl:
         break
+    torch.set_grad_enabled(False)
 
     light_model.training_step(batch, 0)
 

diff --git a/tests/models/efficient_det/lightning/test_validate.py b/tests/models/efficient_det/lightning/test_validate.py
@@ -49,6 +49,7 @@ def test_lightining_efficientdet_finalizes_metrics_on_validation_epoch_end(
 ):
     light_model = light_model_cls(fridge_efficientdet_model, metrics=metrics)
     light_model.convert_raw_predictions = lambda *args: None
+    torch.set_grad_enabled(False)
 
     light_model.validation_epoch_end(None)
 
@@ -65,6 +66,7 @@ def test_lightining_efficientdet_logs_losses_during_validation_step(
     light_model.convert_raw_predictions = lambda *args: None
     light_model.compute_loss = lambda *args: None
     light_model.accumulate_metrics = lambda *args: None
+    torch.set_grad_enabled(False)
 
     light_model.validation_step(batch, 0)
 

diff --git a/tests/models/fastai/unet/lightning/test_test.py b/tests/models/fastai/unet/lightning/test_test.py
@@ -55,6 +55,7 @@ def log(self, key, value, **args):
     for batch in valid_dl:
         break
     light_model = LightModel(model)
+    torch.set_grad_enabled(False)
 
     light_model.test_step(batch, 0)
 

diff --git a/tests/models/fastai/unet/lightning/test_train.py b/tests/models/fastai/unet/lightning/test_train.py
@@ -51,6 +51,7 @@ def configure_optimizers(self):
     expected_loss = random.randint(0, 1000)
     light_model = LightModel(model)
     light_model.compute_loss = lambda *args: expected_loss
+    torch.set_grad_enabled(False)
 
     loss = light_model.training_step(batch, 0)
 
@@ -83,6 +84,7 @@ def log(self, key, value, **args):
     for batch in train_dl:
         break
     light_model = LightModel(model)
+    torch.set_grad_enabled(False)
 
     light_model.training_step(batch, 0)
 

diff --git a/tests/models/fastai/unet/lightning/test_validate.py b/tests/models/fastai/unet/lightning/test_validate.py
@@ -55,6 +55,7 @@ def log(self, key, value, **args):
     for batch in valid_dl:
         break
     light_model = LightModel(model)
+    torch.set_grad_enabled(False)
 
     light_model.validation_step(batch, 0)
 

diff --git a/tests/models/mmdet/test_model.py b/tests/models/mmdet/test_model.py
@@ -89,6 +89,7 @@ def configure_optimizers(self):
         for batch in train_dl:
             batch
             break
+        torch.set_grad_enabled(False)
 
         loss = light_model.training_step(batch, 0)
 
@@ -123,6 +124,7 @@ def log(self, key, value, **args):
         for batch in train_dl:
             batch
             break
+        torch.set_grad_enabled(False)
 
         light_model.training_step(batch, 0)
 
@@ -178,6 +180,7 @@ def log(self, key, value, **args):
         for batch in valid_dl:
             batch
             break
+        torch.set_grad_enabled(False)
 
         light_model.validation_step(batch, 0)
 
@@ -202,6 +205,7 @@ def finalize_metrics(self):
                 self.was_finalize_metrics_called = True
 
         light_model = LitModel(model)
+        torch.set_grad_enabled(False)
 
         light_model.validation_epoch_end(None)
 
@@ -257,6 +261,7 @@ def log(self, key, value, **args):
         for batch in valid_dl:
             batch
             break
+        torch.set_grad_enabled(False)
 
         light_model.test_step(batch, 0)
 
@@ -281,6 +286,7 @@ def finalize_metrics(self):
                 self.was_finalize_metrics_called = True
 
         light_model = LitModel(model)
+        torch.set_grad_enabled(False)
 
         light_model.test_epoch_end(None)
 

diff --git a/tests/models/torchvision_models/faster_rcnn/lightning/test_test.py b/tests/models/torchvision_models/faster_rcnn/lightning/test_test.py
@@ -46,6 +46,7 @@ def test_lightining_faster_rcnn_finalizes_metrics_on_test_epoch_end(
     fridge_faster_rcnn_model, light_model_cls, metrics
 ):
     light_model = light_model_cls(fridge_faster_rcnn_model, metrics=metrics)
+    torch.set_grad_enabled(False)
 
     light_model.test_epoch_end(None)
 
@@ -61,6 +62,7 @@ def test_lightining_faster_rcnn_logs_losses_during_test_step(
         break
     light_model.convert_raw_predictions = lambda **args: None
     light_model.accumulate_metrics = lambda **args: None
+    torch.set_grad_enabled(False)
 
     light_model.test_step(batch, 0)
 

diff --git a/tests/models/torchvision_models/faster_rcnn/lightning/test_train.py b/tests/models/torchvision_models/faster_rcnn/lightning/test_train.py
@@ -48,6 +48,7 @@ def test_lightining_faster_rcnn_training_step_returns_loss(
     light_model.compute_loss = lambda *args: expected_loss
     for batch in train_dl:
         break
+    torch.set_grad_enabled(False)
 
     loss = light_model.training_step(batch, 0)
 
@@ -61,6 +62,7 @@ def test_lightining_faster_rcnn_logs_losses_during_training_step(
     light_model = light_model_cls(fridge_faster_rcnn_model, metrics=None)
     for batch in train_dl:
         break
+    torch.set_grad_enabled(False)
 
     light_model.training_step(batch, 0)
 

diff --git a/tests/models/torchvision_models/faster_rcnn/lightning/test_validate.py b/tests/models/torchvision_models/faster_rcnn/lightning/test_validate.py
@@ -46,6 +46,7 @@ def test_lightining_faster_rcnn_finalizes_metrics_on_validation_epoch_end(
     fridge_faster_rcnn_model, light_model_cls, metrics
 ):
     light_model = light_model_cls(fridge_faster_rcnn_model, metrics=metrics)
+    torch.set_grad_enabled(False)
 
     light_model.validation_epoch_end(None)
 
@@ -61,6 +62,7 @@ def test_lightining_faster_rcnn_logs_losses_during_validation_step(
         break
     light_model.convert_raw_predictions = lambda **args: None
     light_model.accumulate_metrics = lambda **args: None
+    torch.set_grad_enabled(False)
 
     light_model.validation_step(batch, 0)
 

diff --git a/tests/models/torchvision_models/keypoints_rcnn/lightning/test_test.py b/tests/models/torchvision_models/keypoints_rcnn/lightning/test_test.py
@@ -42,6 +42,7 @@ def test_lightining_keypoints_rcnn_test(ochuman_keypoints_dls, light_model_cls):
 def test_lightining_keypoints_finalizes_metrics_on_test_epoch_end(light_model_cls):
     model = keypoint_rcnn.model(num_keypoints=19)
     light_model = light_model_cls(model)
+    torch.set_grad_enabled(False)
 
     light_model.test_epoch_end(None)
 
@@ -58,6 +59,7 @@ def test_lightining_keypoints_rcnn_logs_losses_during_test_step(
         break
     light_model.convert_raw_predictions = lambda **args: None
     light_model.accumulate_metrics = lambda **args: None
+    torch.set_grad_enabled(False)
 
     light_model.test_step(batch, 0)
 

diff --git a/tests/models/torchvision_models/keypoints_rcnn/lightning/test_train.py b/tests/models/torchvision_models/keypoints_rcnn/lightning/test_train.py
@@ -46,6 +46,7 @@ def test_lightining_keypoints_rcnn_training_step_returns_loss(
     light_model.compute_loss = lambda *args: expected_loss
     for batch in train_dl:
         break
+    torch.set_grad_enabled(False)
 
     loss = light_model.training_step(batch, 0)
 
@@ -60,6 +61,7 @@ def test_lightining_keypoints_rcnn_logs_losses_during_training_step(
     light_model = light_model_cls(model)
     for batch in train_dl:
         break
+    torch.set_grad_enabled(False)
 
     light_model.training_step(batch, 0)
 

diff --git a/tests/models/torchvision_models/keypoints_rcnn/lightning/test_validate.py b/tests/models/torchvision_models/keypoints_rcnn/lightning/test_validate.py
@@ -44,6 +44,7 @@ def test_lightining_keypoints_finalizes_metrics_on_validation_epoch_end(
 ):
     model = keypoint_rcnn.model(num_keypoints=19)
     light_model = light_model_cls(model)
+    torch.set_grad_enabled(False)
 
     light_model.validation_epoch_end(None)
 
@@ -60,6 +61,7 @@ def test_lightining_keypoints_rcnn_logs_losses_during_validation_step(
         break
     light_model.convert_raw_predictions = lambda **args: None
     light_model.accumulate_metrics = lambda **args: None
+    torch.set_grad_enabled(False)
 
     light_model.validation_step(batch, 0)
 

diff --git a/tests/models/torchvision_models/retinanet/lightning/test_test.py b/tests/models/torchvision_models/retinanet/lightning/test_test.py
@@ -46,6 +46,7 @@ def test_lightining_retinanet_finalizes_metrics_on_test_epoch_end(
     fridge_retinanet_model, light_model_cls, metrics
 ):
     light_model = light_model_cls(fridge_retinanet_model, metrics=metrics)
+    torch.set_grad_enabled(False)
 
     light_model.test_epoch_end(None)
 
@@ -61,6 +62,7 @@ def test_lightining_retinanet_logs_losses_during_test_step(
         break
     light_model.convert_raw_predictions = lambda **args: None
     light_model.accumulate_metrics = lambda **args: None
+    torch.set_grad_enabled(False)
 
     light_model.test_step(batch, 0)
 

diff --git a/tests/models/torchvision_models/retinanet/lightning/test_train.py b/tests/models/torchvision_models/retinanet/lightning/test_train.py
@@ -46,6 +46,7 @@ def test_lightining_retinanet_training_step_returns_loss(
     light_model.compute_loss = lambda *args: expected_loss
     for batch in train_dl:
         break
+    torch.set_grad_enabled(False)
 
     loss = light_model.training_step(batch, 0)
 
@@ -59,6 +60,7 @@ def test_lightining_retinanet_logs_losses_during_training_step(
     light_model = light_model_cls(fridge_retinanet_model)
     for batch in train_dl:
         break
+    torch.set_grad_enabled(False)
 
     light_model.training_step(batch, 0)
 

diff --git a/tests/models/torchvision_models/retinanet/lightning/test_validate.py b/tests/models/torchvision_models/retinanet/lightning/test_validate.py
@@ -46,6 +46,7 @@ def test_lightining_retinanet_finalizes_metrics_on_validation_epoch_end(
     fridge_retinanet_model, light_model_cls, metrics
 ):
     light_model = light_model_cls(fridge_retinanet_model, metrics=metrics)
+    torch.set_grad_enabled(False)
 
     light_model.validation_epoch_end(None)
 
@@ -61,6 +62,7 @@ def test_lightining_retinanet_logs_losses_during_validation_step(
         break
     light_model.convert_raw_predictions = lambda **args: None
     light_model.accumulate_metrics = lambda **args: None
+    torch.set_grad_enabled(False)
 
     light_model.validation_step(batch, 0)
 

diff --git a/tests/models/ultralytics/yolov5/lightning/test_test.py b/tests/models/ultralytics/yolov5/lightning/test_test.py
@@ -49,6 +49,7 @@ def test_lightning_yolo_finalizes_metrics_on_test_epoch_end(backbone, light_mode
     )
     metrics = [COCOMetric(metric_type=COCOMetricType.bbox)]
     light_model = light_model_cls(model, metrics=metrics)
+    torch.set_grad_enabled(False)
 
     light_model.test_epoch_end(None)
 
@@ -88,6 +89,7 @@ def log(self, key, value, **args):
     for batch in valid_dl:
         batch
         break
+    torch.set_grad_enabled(False)
 
     light_model.test_step(batch, 0)
 

diff --git a/tests/models/ultralytics/yolov5/lightning/test_train.py b/tests/models/ultralytics/yolov5/lightning/test_train.py
@@ -56,6 +56,7 @@ def configure_optimizers(self):
     for batch in train_dl:
         batch
         break
+    torch.set_grad_enabled(False)
 
     loss = light_model.training_step(batch, 0)
 
@@ -94,6 +95,7 @@ def log(self, key, value, **args):
     for batch in train_dl:
         batch
         break
+    torch.set_grad_enabled(False)
 
     light_model.training_step(batch, 0)
 

diff --git a/tests/models/ultralytics/yolov5/lightning/test_validate.py b/tests/models/ultralytics/yolov5/lightning/test_validate.py
@@ -51,6 +51,7 @@ def test_lightning_yolo_finalizes_metrics_on_validation_epoch_end(
     )
     metrics = [COCOMetric(metric_type=COCOMetricType.bbox)]
     light_model = light_model_cls(model, metrics=metrics)
+    torch.set_grad_enabled(False)
 
     light_model.validation_epoch_end(None)
 
@@ -90,6 +91,7 @@ def log(self, key, value, **args):
     for batch in valid_dl:
         batch
         break
+    torch.set_grad_enabled(False)
 
     light_model.validation_step(batch, 0)