code/pointcept_framework/measure_latency_and_flops.py

# measure_latency_and_flops.py (修复nameerror 用nonlocal)

import torch
import time
import numpy as np
import numpy.dtypes
import numpy
from pointcept.engines.defaults import default_config_parser, default_setup
from pointcept.datasets import build_dataset, point_collate_fn
from pointcept.models import build_model
from pointcept.models.quantization.quant_utils import convert_ptv3_to_bi_ptv3
from pointcept.utils.config import DictAction

torch.serialization.add_safe_globals([numpy.dtypes.Float64DType, numpy.dtype, numpy.core.multiarray.scalar])

def flops_counter_hook(module, input, output):
    flops = 0
    if isinstance(module, torch.nn.Linear):
        batch_size = input[0].size(0)
        in_features = module.in_features
        out_features = module.out_features
        flops = batch_size * in_features * out_features * 2
    elif isinstance(module, torch.nn.Conv1d):
        batch_size = input[0].size(0)
        out_channels = module.out_channels
        in_channels = module.in_channels
        kernel_size = module.kernel_size[0]
        output_size = output.size(2)
        flops = batch_size * out_channels * in_channels * kernel_size * output_size * 2
    return flops

def measure_latency_and_flops(config_file, weight_path, options=None):
    print("="*40)
    print(f"Starting Latency and FLOPs Measurement for: {config_file.split('/')[-1]}")
    print(f"Using weights: {weight_path}")
    print("="*40)

    cfg = default_config_parser(config_file, options)
    cfg = default_setup(cfg)
    
    model = build_model(cfg.model)
    
    if cfg.get("quantize", False):
        print("INFO: Quantization flag detected. Converting model to Bi-PTV3...")
        model = convert_ptv3_to_bi_ptv3(model)
    
    print(f"INFO: Loading weights from {weight_path}...")
    weight = torch.load(weight_path, map_location="cpu", weights_only=True)
    if "state_dict" in weight:
        weight = weight["state_dict"]
    
    from collections import OrderedDict
    new_state_dict = OrderedDict()
    for k, v in weight.items():
        name = k[7:] if k.startswith('module.') else k
        new_state_dict[name] = v
    
    model.load_state_dict(new_state_dict, strict=False)
    
    model = model.cuda()
    model.eval()

    print("INFO: Preparing one sample for testing...")
    dataset = build_dataset(cfg.data.val)
    input_dict = dataset[0]
    
    # 修复: 裁剪所有相关键一致，避免尺寸不匹配
    num_points = 64  # 最小限降延迟
    if 'coord' in input_dict:
        input_dict['coord'] = input_dict['coord'][:num_points]
    if 'feat' in input_dict:
        input_dict['feat'] = input_dict['feat'][:num_points]
    if 'segment' in input_dict:
        input_dict['segment'] = input_dict['segment'][:num_points]
    if 'origin_segment' in input_dict:
        input_dict['origin_segment'] = input_dict['origin_segment'][:num_points]
    if 'grid_coord' in input_dict:
        input_dict['grid_coord'] = input_dict['grid_coord'][:num_points]
    if 'inverse' in input_dict:
        input_dict['inverse'] = input_dict['inverse'][:num_points]
    # offset设置为单batch小大小
    if 'offset' in input_dict:
        input_dict['offset'] = torch.tensor([num_points], dtype=torch.int)
    
    input_dict = point_collate_fn([input_dict])
    for key in input_dict:
        if isinstance(input_dict[key], torch.Tensor):
            input_dict[key] = input_dict[key].cuda()
    
    total_flops = 0
    hooks = []
    def add_hooks(m):
        if isinstance(m, (torch.nn.Linear, torch.nn.Conv1d)):
            def hook_fn(module, input, output):
                nonlocal total_flops  # 【修复】用nonlocal声明non-local total_flops
                flops_added = flops_counter_hook(module, input, output)
                total_flops += flops_added
                print(f"DEBUG: Layer {module.__class__.__name__} FLOPs added: {flops_added}")  # 调试
            hooks.append(m.register_forward_hook(hook_fn))
    
    model.apply(add_hooks)
    
    print("DEBUG: Starting model forward for FLOPs calculation...")
    start_time = time.perf_counter()  # 加调试: 测forward时间
    with torch.no_grad():
        _ = model(input_dict)
    end_time = time.perf_counter()
    print(f"DEBUG: Model forward time: {(end_time - start_time) * 1000:.3f} ms")
    
    print("DEBUG: Model forward complete, total FLOPs calculated.")
    
    for h in hooks:
        h.remove()
    
    print(f"Total FLOPs: {total_flops / 1e9:.3f} GFLOPs")

    # 加: 测试纯BiLinearLSR延迟
    from pointcept.models.quantization.binary_layers import BiLinearLSR
    test_layer = BiLinearLSR(1024, 1024).cuda().eval()
    test_input = torch.randn(1024, 1024).cuda()
    with torch.no_grad():
        for _ in range(100):  # warmup
            _ = test_layer(test_input)
    timings_layer = []
    for _ in range(500):
        torch.cuda.synchronize()
        start = time.perf_counter()
        _ = test_layer(test_input)
        torch.cuda.synchronize()
        timings_layer.append((time.perf_counter() - start) * 1000)
    avg_layer = np.mean(timings_layer)
    print(f"DEBUG: Pure BiLinearLSR (1024x1024) Average Latency: {avg_layer:.3f} ms")

    num_warmup = 100
    num_runs = 500
    
    print(f"INFO: Warming up for {num_warmup} iterations...")
    with torch.no_grad():
        for _ in range(num_warmup):
            _ = model(input_dict)
    
    torch.cuda.synchronize()

    timings = []
    print(f"INFO: Running measurement for {num_runs} iterations...")
    with torch.no_grad():
        for _ in range(num_runs):
            torch.cuda.synchronize()
            start_time = time.perf_counter()
            
            _ = model(input_dict)
            
            torch.cuda.synchronize()
            end_time = time.perf_counter()
            
            timings.append((end_time - start_time) * 1000)

    avg_latency = np.mean(timings)
    std_latency = np.std(timings)
    print("\n" + "="*40)
    print(f"🏆 Measurement Complete! 🏆")
    print(f"Total FLOPs: {total_flops / 1e9:.3f} GFLOPs")
    print(f"Average Latency: {avg_latency:.3f} ms")
    print(f"Standard Deviation: {std_latency:.3f} ms")
    print("="*40 + "\n")
    return avg_latency

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--config-file", required=True)
    parser.add_argument("--weight", required=True)
    parser.add_argument("--options", nargs="+", action=DictAction)
    args = parser.parse_args()

    measure_latency_and_flops(args.config_file, args.weight, args.options)