Function estimate_size

utils/interp.py:45–48

(model)

def model_size_estimation(module):
    def estimate_size(model):
        param_size = sum(p.nelement() * p.element_size() for p in model.parameters())
        buffer_size = sum(b.nelement() * b.element_size() for b in model.buffers())
        return param_size + buffer_size

    total_size = estimate_size(module)
    output = [f'Total model size: {total_size / 1024**2:.2f} MB', '---------------------------']