Skip to content

Add Bytes and KB Units for small PyTorch models #315

Open
@halsimov

Description

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
==========================================================================================

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions