Deprecate the model_size property on the LightningModule

[ananthsub]

🚀 Feature

We are auditing the Lightning components and APIs to assess opportunities for improvements:

• https://docs.google.com/document/d/1xHU7-iQSpp9KJTjI3As2EM0mfNHHr37WZYpDpwLkivA/edit#
• Review Lightning architecture & API #7740 (comment)

One item that came up on the LightningModule API was the ambiguity of the model_size property: https://github.com/PyTorchLightning/pytorch-lightning/blob/1ad1a89c09e6e8857f4fe680e23644a5013011d0/pytorch_lightning/core/lightning.py#L2073-L2084

The value this property provides is ambiguous, and its implementation of dumping the module to a file to check the size is suboptimal. To reflect this, we'd like to deprecate the property off the core API. If this is still needed, this could be a utility function which accepts an nn.Module and returns the size.

@awaelchli @Borda

[tchaton]

Hryt @ananthsub,

I think it is a good remark and it won t scale well with larger model.
Mind doing the depreciation ?

Best,
T.C

[awaelchli]

Hi @ananthsub
This was a contribution by the community. I vote for deprecation in favor of a utility function :)

[calebrob6]

I noticed the sub-optimal implementation of model_size and came up with two (maybe better) options:

@property
def model_size(self):
    # Uses a BytesIO buffer instead of writing to disk (this will cost RAM so may not be desirable)
    # See https://stackoverflow.com/questions/26827055/python-how-to-get-bytesio-allocated-memory-length/26827410
    with io.BytesIO() as f:
        torch.save(self.state_dict(), f)
        f.seek(0, io.SEEK_END)
        size_mb = f.tell() / 1e6
    return size_mb

@property
def model_size(self):
    # This is a naive way to get the size of the tensors in the `state_dict()` per discussion here https://discuss.pytorch.org/t/how-to-know-the-memory-allocated-for-a-tensor-on-gpu/28537        
    # Better way to do it https://github.com/Stonesjtu/pytorch_memlab/blob/master/pytorch_memlab/mem_reporter.py#L17
    size_mb = 0
    for k, v in self.state_dict().items():
        size_mb += v.element_size() * v.nelement()
    size_mb /= 1e6
    return size_mb

I would be happy to do this in a utility function!

[ananthsub]

I would be happy to do this in a utility function!

Thanks @calebrob6 ! #8495 will deprecate this off the core LightningModule API. A separate function under https://github.com/PyTorchLightning/pytorch-lightning/tree/master/pytorch_lightning/utilities that accepts an nn.Module could be useful in parallel.

If you have a chance, feel free to contribute this 😄

[carmocca]

This was a contribution by the community. I vote for deprecation in favor of a utility function :)

Not really, this was added by @Borda in

b434c47#diff-21dfd052fd2cf070881d88cd38b207ebc1764be99c3d289208f36a6a630378a5R1769-R1776
(PR #5706)

as we already have a model_size "utility" in:
https://github.com/PyTorchLightning/pytorch-lightning/blob/4c79b3a5b343866217784c66d122819c59a92c1d/pytorch_lightning/core/memory.py#L261-L264

So the deprecation of the property in the LightningModule should also address its replacement in test_quantization.py

[calebrob6]

I added a simple utility function + tests in #8500 before I saw the above. @carmocca are you saying that the model size functionality should belong in the ModelSummary class?

[carmocca]

are you saying that the model size functionality should belong in the ModelSummary class?

Yes, as it is already supposed to compute model_size