summarize total size of model params in bytes

pytorch_lightning/core/memory.py

@@ -159,6 +159,7 @@ class ModelSummary(object):
         132 K     Trainable params
         0         Non-trainable params
         132 K     Total params
+        0.506     Total estimated model params size (MB)
         >>> ModelSummary(model, mode='full')  # doctest: +NORMALIZE_WHITESPACE
           | Name  | Type        | Params | In sizes  | Out sizes
         --------------------------------------------------------------
@@ -169,6 +170,7 @@ class ModelSummary(object):
         132 K     Trainable params
         0         Non-trainable params
         132 K     Total params
+        0.506     Total estimated model params size (MB)
     """
 
     MODE_TOP = "top"
@@ -180,6 +182,7 @@ def __init__(self, model, mode: str = MODE_DEFAULT):
         self._model = model
         self._mode = mode
         self._layer_summary = self.summarize()
+        self._precision_megabytes = (self._model.precision / 8.0) / (1024 ** 2.0) # 1 byte -> 8 bits)
 
     @property
     def named_modules(self) -> List[Tuple[str, nn.Module]]:
@@ -213,6 +216,29 @@ def out_sizes(self) -> List:
     def param_nums(self) -> List[int]:
         return [layer.num_parameters for layer in self._layer_summary.values()]
 
+    @property
+    def total_parameters(self) -> int:
+        return sum(p.numel() for p in self._model.parameters())
+
+    @property
+    def trainable_parameters(self) -> int:
+        return sum(p.numel() for p in self._model.parameters() if p.requires_grad)
+
+    @property
+    def model_size(self) -> float:
+        return self._model_size()
+    def _model_size(self) -> float:
+        """
+        Estimates total model size i.e total params size in MBs
+        total params size gives model size in accounting total model params.
+
+        NOTE: Currently only Supported total params size.
+
+        Returns:
+            Total estimated model size(MB).
+        """
+        return self.total_parameters * self._precision_megabytes
+
     def summarize(self) -> Dict[str, LayerSummary]:
         summary = OrderedDict((name, LayerSummary(module)) for name, module in self.named_modules)
         if self._model.example_input_array is not None:
@@ -259,11 +285,11 @@ def __str__(self):
         if self._model.example_input_array is not None:
             arrays.append(["In sizes", self.in_sizes])
             arrays.append(["Out sizes", self.out_sizes])
+        total_parameters = self.total_parameters
+        trainable_parameters = self.trainable_parameters
+        model_size = self.model_size
 
-        trainable_parameters = sum(p.numel() for p in self._model.parameters() if p.requires_grad)
-        total_parameters = sum(p.numel() for p in self._model.parameters())
-
-        return _format_summary_table(total_parameters, trainable_parameters, *arrays)
+        return _format_summary_table(total_parameters, trainable_parameters, model_size, *arrays)
 
     def __repr__(self):
         return str(self)
@@ -280,7 +306,7 @@ def parse_batch_shape(batch: Any) -> Union[str, List]:
     return UNKNOWN_SIZE
 
 
-def _format_summary_table(total_parameters: int, trainable_parameters: int, *cols) -> str:
+def _format_summary_table(total_parameters: int, trainable_parameters: int, model_size: float, *cols) -> str:
     """
     Takes in a number of arrays, each specifying a column in
     the summary table, and combines them all into one big
@@ -316,6 +342,8 @@ def _format_summary_table(total_parameters: int, trainable_parameters: int, *col
     summary += "Non-trainable params"
     summary += "\n" + s.format(get_human_readable_count(total_parameters), 10)
     summary += "Total params"
+    summary += "\n" + s.format(get_formatted_model_size(model_size), 10)
+    summary += "Total estimated model params size (MB)"
 
     return summary
 
@@ -372,6 +400,8 @@ def get_gpu_memory_map() -> Dict[str, int]:
     }
     return gpu_memory_map
 
+def get_formatted_model_size(total_model_size: float) -> float:
+    return f"{total_model_size:,.3f}"
 
 def get_human_readable_count(number: int) -> str:
     """
@@ -411,3 +441,6 @@ def get_human_readable_count(number: int) -> str:
         return f"{int(number):,d} {labels[index]}"
     else:
         return f"{number:,.1f} {labels[index]}"
+if __name__ == "__main__":
+    import doctest
+    doctest.testmod()

tests/core/test_memory.py

@@ -33,6 +33,21 @@ def forward(self, *args, **kwargs):
         return {'loss': self.parameter.sum()}
 
 
+class PreCalculatedModel(LightningModule):
+    """ A module with precalculated total params size in MB. """
+
+    def __init__(self):
+        super().__init__()
+        self.layer1 = nn.Linear(10, 100)
+        self.layer2 = nn.Linear(100, 2)
+        self.pre_calculated_model_size = 0.005
+
+    def forward(self, x):
+        x = self.layer1(x)
+        x = self.layer2(x)
+        return x
+
+
 class UnorderedModel(LightningModule):
     """ A model in which the layers not defined in order of execution """
 
@@ -247,3 +262,25 @@ def forward(self, *args, **kwargs):
     model.example_input_array = example_input
     summary = model.summarize(mode=mode)
     assert summary.in_sizes == [expected_size]
+
+@pytest.mark.parametrize(['mode'], [
+    pytest.param(ModelSummary.MODE_FULL),
+    pytest.param(ModelSummary.MODE_TOP),
+])
+def test_model_size(mode):
+    """ Test that model size is calculated correctly. """
+    model = PreCalculatedModel()
+    summary = model.summarize(mode=mode)
+    pre_calculated_model_size = torch.tensor(model.pre_calculated_model_size)
+    model_size = torch.tensor(summary.model_size)
+    assert torch.isclose(model_size, pre_calculated_model_size, atol=1e-4)
+
+@pytest.mark.parametrize(['mode'], [
+    pytest.param(ModelSummary.MODE_FULL),
+    pytest.param(ModelSummary.MODE_TOP),
+])
+def test_empty_model_size(mode):
+    """ Test that empty model size is zero. """
+    model = EmptyModule()
+    summary = model.summarize(mode=mode)
+    assert 0.0 == summary.model_size