Function count_params

models/utils.py:211–215

(model, verbose=False)

    model.total_ops += torch.DoubleTensor([matmul_ops])

def count_params(model, verbose=False):
    total_params = sum(p.numel() for p in model.parameters())
    if verbose:
        print(f"{model.__class__.__name__} has {total_params * 1.e-6:.2f} M params.")
    return total_params