Deprecate LightningModule.model_size

CHANGELOG.md

@@ -39,7 +39,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Deprecated
 
--
+- Deprecated `LightningModule.model_size` ([#8343](https://github.com/PyTorchLightning/pytorch-lightning/pull/8343))
 
 
 -

pytorch_lightning/core/lightning.py

@@ -1959,7 +1959,11 @@ def model_size(self) -> float:
         The model's size in megabytes. The computation includes everything in the
         :meth:`~torch.nn.Module.state_dict`, i.e., by default the parameteters and buffers.
         """
-        # todo: think about better way without need to dump model to drive
+        rank_zero_deprecation(
+            "The `LightningModule.model_size` property was deprecated in v1.5 and will be removed in v1.7. Please "
+            "use the `utilities.memory.get_model_size_mb` method under utilities/memory.py"
+        )
+
         tmp_name = f"{uuid.uuid4().hex}.pt"
         torch.save(self.state_dict(), tmp_name)
         size_mb = os.path.getsize(tmp_name) / 1e6

pytorch_lightning/utilities/memory.py

@@ -13,8 +13,11 @@
 # limitations under the License.
 
 import gc
+import os
+import uuid
 
 import torch
+from torch.nn import Module
 
 
 def recursive_detach(in_dict: dict, to_cpu: bool = False) -> dict:
@@ -87,3 +90,18 @@ def garbage_collection_cuda():
         if not is_oom_error(exception):
             # Only handle OOM errors
             raise
+
+
+def get_model_size_mb(model: Module) -> float:
+    """
+    Calculates the size of a Module in megabytes by saving the model to a temporary file
+    and reading in the size.
+    Returns:
+        Number of megabytes in the parameters of the input module
+    """
+    # TODO: Implement a method without needing to download the model
+    tmp_name = f"{uuid.uuid4().hex}.pt"
+    torch.save(model.state_dict(), tmp_name)
+    size_mb = os.path.getsize(tmp_name) / 1e6
+    os.remove(tmp_name)
+    return size_mb

tests/callbacks/test_quantization.py

@@ -21,6 +21,7 @@
 from pytorch_lightning import seed_everything, Trainer
 from pytorch_lightning.callbacks import QuantizationAwareTraining
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
+from pytorch_lightning.utilities.memory import get_model_size_mb
 from tests.helpers.datamodules import RegressDataModule
 from tests.helpers.runif import RunIf
 from tests.helpers.simple_models import RegressionModel
@@ -40,7 +41,7 @@ def test_quantization(tmpdir, observe: str, fuse: bool, convert: bool):
 
     trainer = Trainer(**trainer_args)
     trainer.fit(model, datamodule=dm)
-    org_size = model.model_size
+    org_size = get_model_size_mb(model)
     org_score = torch.mean(torch.tensor([mean_relative_error(model(x), y) for x, y in dm.test_dataloader()]))
 
     fusing_layers = [(f"layer_{i}", f"layer_{i}a") for i in range(3)] if fuse else None
@@ -62,7 +63,7 @@ def test_quantization(tmpdir, observe: str, fuse: bool, convert: bool):
         qmodel.eval()
         torch.quantization.convert(qmodel, inplace=True)
 
-    quant_size = qmodel.model_size
+    quant_size = get_model_size_mb(qmodel)
     # test that the trained model is smaller then initial
     size_ratio = quant_size / org_size
     assert size_ratio < 0.65

tests/deprecated_api/test_remove_1-7.py

@@ -0,0 +1,26 @@
+# Copyright The PyTorch Lightning team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+""" Test deprecated functionality which will be removed in v1.7.0 """
+import pytest
+
+from tests.helpers import BoringModel
+
+
+def test_v1_7_0_deprecated_model_size():
+    model = BoringModel()
+    with pytest.deprecated_call(
+        match="The `LightningModule.model_size` property was deprecated in v1.5 and will be removed in v1.7. "
+        "Please use the `utilities.memory.get_model_size_mb` method under utilities/memory.py"
+    ):
+        _ = model.model_size

tests/utilities/test_memory.py

@@ -11,9 +11,13 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import math
+
 import torch
+import torch.nn as nn
 
-from pytorch_lightning.utilities.memory import recursive_detach
+from pytorch_lightning.utilities.memory import get_model_size_mb, recursive_detach
+from tests.helpers import BoringModel
 
 
 def test_recursive_detach():
@@ -28,3 +32,24 @@ def test_recursive_detach():
     assert y["foo"].device.type == "cpu"
     assert y["bar"]["baz"].device.type == "cpu"
     assert not y["bar"]["baz"].requires_grad
+
+
+def test_get_model_size_mb():
+    model = BoringModel()
+
+    size_bytes = get_model_size_mb(model)
+
+    # Size will be python version dependent.
+    assert math.isclose(size_bytes, 0.001319, rel_tol=0.1)
+
+
+def test_get_sparse_model_size_mb():
+    class BoringSparseModel(BoringModel):
+        def __init__(self):
+            super().__init__()
+            self.layer = nn.Parameter(torch.ones(32).to_sparse())
+
+    model = BoringSparseModel()
+    size_bytes = get_model_size_mb(model)
+
+    assert math.isclose(size_bytes, 0.001511, rel_tol=0.1)