Function module_stats_hook

imperative/python/megengine/utils/module_stats.py:483–523

(module, inputs, outputs, name="")

    recorded_parameters = set()

    def module_stats_hook(module, inputs, outputs, name=""):
        class_name = str(module.__class__).split(".")[-1].split("'")[0]
        if cal_flops:
            flops_stats = get_op_stats(module, inputs, outputs)
            if flops_stats is not None:
                flops_stats["name"] = name
                flops_stats["class_name"] = class_name
                flops.append(flops_stats)

        if cal_params:
            if (
                hasattr(module, "weight")
                and (module.weight is not None)
                and module.weight not in recorded_parameters
            ):
                w = module.weight
                param_stats = get_param_stats(w)
                param_stats["name"] = name + "-w"
                params.append(param_stats)
                recorded_parameters.add(w)

            if (
                hasattr(module, "bias")
                and module.bias is not None
                and module.bias not in recorded_parameters
            ):
                b = module.bias
                param_stats = get_param_stats(b)
                param_stats["name"] = name + "-b"
                params.append(param_stats)
                recorded_parameters.add(b)

        if cal_activations:
            if not isinstance(outputs, (tuple, list)):
                output = outputs
            else:
                output = outputs[0]
            activation_stats = get_activation_stats(output, has_inputs)
            activation_stats["name"] = name
            activation_stats["class_name"] = class_name
            activations.append(activation_stats)

    params = []
    flops = []