调用函数

检查权重参数

def assert_model_params_finite(module: nn.Module) -> None:
    bad_params = []
    for name, param in module.named_parameters():
        nan_count = torch.isnan(param).sum().item()
        inf_count = torch.isinf(param).sum().item()
        if nan_count or inf_count:
            bad_params.append(
                f"{name}: shape={tuple(param.shape)}, nan={nan_count}, inf={inf_count}"
            )
    if bad_params:
        details = "\n".join(bad_params)
        raise ValueError(
            "Model parameters contain NaN/Inf before quantization.\n"
            "Quantization observer failure is a downstream symptom.\n"
            f"{details}"
        )

更进一步检查bn层有没有nan和inf值

def check_bn_running_stats(module: nn.Module) -> None:
    bad_bn_stats = []
    bn_types = (
        nn.BatchNorm1d,
        nn.BatchNorm2d,
        nn.BatchNorm3d,
        nn.SyncBatchNorm,
    )
    for name, submodule in module.named_modules():
        if not isinstance(submodule, bn_types):
            continue
        for stat_name in ("running_mean", "running_var"):
            stat = getattr(submodule, stat_name, None)
            if stat is None:
                continue
            nan_count = torch.isnan(stat).sum().item()
            inf_count = torch.isinf(stat).sum().item()
            if nan_count or inf_count:
                finite_stat = stat[torch.isfinite(stat)]
                if finite_stat.numel():
                    min_val = finite_stat.min().item()
                    max_val = finite_stat.max().item()
                else:
                    min_val = "N/A"
                    max_val = "N/A"
                bad_bn_stats.append(
                    f"{name}.{stat_name}: shape={tuple(stat.shape)}, "
                    f"nan={nan_count}, inf={inf_count}, min={min_val}, max={max_val}"
                )
    if bad_bn_stats:
        print("Abnormal BatchNorm running stats found:")
        for item in bad_bn_stats:
            print(item)
    else:
        print("All BatchNorm running_mean/running_var are finite.")

使用

    model = YOLO(pt_path)
    # assert_model_params_finite(model)
    # check_bn_running_stats(model)