# 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)