torchinfo icon indicating copy to clipboard operation
torchinfo copied to clipboard

Add Bytes and KB Units for small PyTorch models

Open halsimov opened this issue 1 year ago • 1 comments

version: torchinfo 1.8.0 pyhd8ed1ab_0

Describe the solution you'd like

Add the Units alongside MB and GB and choose unit acoording to the best representation one can get.

In the following toy example having MB as unit is not of any help :

==========================================================================================
Layer (type:depth-idx)                   Output Shape              Param #
==========================================================================================
LRBasedClassifier                        [10, 5]                   --
├─Linear: 1-1                            [10, 5]                   125
==========================================================================================
Total params: 125
Trainable params: 125
Non-trainable params: 0
Total mult-adds (Units.MEGABYTES): 0.00
==========================================================================================
Input size (MB): 0.00
Forward/backward pass size (MB): 0.00
Params size (MB): 0.00
Estimated Total Size (MB): 0.00
==========================================================================================

halsimov avatar Jul 16 '24 07:07 halsimov

Good idea, PRs implementing this behavior are welcome! Based on the output value, it should choose the most appropriate unit automatically.

TylerYep avatar Jul 16 '24 20:07 TylerYep