Function get_model_param_count

src/transformers/trainer_pt_utils.py:1009–1023

Calculate model's total param count. If trainable_only is True then count only those requiring grads.

(model, trainable_only=False)

def get_model_param_count(model, trainable_only=False):
    """
    Calculate model's total param count. If trainable_only is True then count only those requiring grads.
    """
    if is_deepspeed_zero3_enabled():

        def numel(p):
            return p.ds_numel if hasattr(p, "ds_numel") else p.numel()

    else:

        def numel(p):
            return p.numel()

    return sum(numel(p) for p in model.parameters() if not trainable_only or p.requires_grad)


def get_parameter_names(model, forbidden_layer_types, forbidden_layer_names=None):