Function param_stat_context

imperative/python/megengine/utils/module_stats.py:543–574

()

    @contextmanager
    def param_stat_context():
        def wrapper(fun):
            @functools.wraps(fun)
            def param_access_record(module, item):
                member = fun(module, item)
                if (
                    item in ["weight", "bias"]
                    and member is not None
                    and member not in recorded_parameters
                ):
                    name = module_to_name[module]
                    if item == "weight":
                        suffix = "-w"
                    elif item == "bias":
                        suffix = "-b"

                    param_name = name + suffix
                    param_stats = get_param_stats(member)
                    param_stats["name"] = param_name
                    params.append(param_stats)
                    recorded_parameters.add(member)

                return member

            return param_access_record

        origin_get_attr = object.__getattribute__
        try:
            M.Module.__getattribute__ = wrapper(origin_get_attr)
            yield
        finally:
            M.Module.__getattribute__ = origin_get_attr

    with set_module_mode_safe(model, training=False) as model, param_stat_context():
        model(*inputs)