Getting model size

[jakedves]

I have the following code from a PyTorch forum that shows us how to compute the model size, however I don't think it takes into account the false-quantisation.

Is there a known way to adapt this code to accurately produce what would be the final model size of a quantised network?

def model_size(model) -> float:
    param_size = 0

    for param in model.parameters():
        param_size += param.nelement() * param.element_size()

    buffer_size = 0
    for buffer in model.buffers():
        buffer_size += buffer.nelement() * buffer.element_size()

    size_all_mb = (param_size + buffer_size) / 1024 ** 2
    return size_all_mb

Thanks!

[Giuseppe5]

Are you sure it does not work for quantization?
All our parameters are either buffers or parameters, since they are exported during the call to state_dict, so they should also be accounted for when iterating through .parameters() and .buffers().

[jakedves]

Hey @Giuseppe5, yes I'm pretty sure, unless I'm misunderstanding something. This script produces the same answer each time:

from brevitas_examples.super_resolution.models import QuantESPCN
from metrics import model_size

for bit_width in [1, 2, 4, 8, 16]:
    model = QuantESPCN(weight_bit_width=bit_width, act_bit_width=bit_width)
    print(model_size(model))
    # 0.2621574401855469
    # 0.2621574401855469
    # 0.2621574401855469
    # 0.2621574401855469
    # 0.2621574401855469

What I would want is a size that takes into account the bit_width that will be exported

Thanks

[Giuseppe5]

Oh I see what you are trying to do.

Since everything is in fake-quantization, it means that the effective data type we use to store the data is always fp32, even though the bitwidth might be different as per your example.

To have an accurate estimation, you would need a custom functions that takes in account the shape of each parameter and the setting that it was used to quantize it.

If you want to be extra precise, you would also need to account for scales (usually stored in fp32/fp16/bf16) and zero point (intX if quantized, otherwise fp32, fp16, bf16).

Unfortunately we do not have such a function at the moment as far as I am aware.