您方便分享一下推理用的数据和onnx模型吗？可能是这个仓库的代码逻辑有问题 : - )

这个网盘里面有我用的pt文件，转出来onnx文件，以及测试图片：https://drive.google.com/drive/folders/1l7ArS2b5TskT4SgPq5duXmyrc7BbOwDq?usp=drive_link

我的导出onnx指令是：
python ./deploy/export.py
--config /home/uestc/nvs/OpenStereo/cfgs/lightstereo/lightstereo_l_kitti.yaml
--weights /home/uestc/nvs/OpenStereo/checkpoints/stereo_anything/StereoAnything-LightStereo_L.pt
--device 0
--imgsz 1280 736
--include onnx
然后测试onnx的代码如下：

import sys
import os
import argparse
import numpy as np
import torch
import onnxruntime
import time
import matplotlib.pyplot as plt
from easydict import EasyDict
from PIL import Image
from pathlib import Path

sys.path.insert(0, './')
from stereo.utils import common_utils
from stereo.modeling import build_trainer
from stereo.utils.disp_color import disp_to_color
from stereo.datasets.dataset_template import build_transform_by_cfg

def parse_config():
    parser = argparse.ArgumentParser(description='ONNX vs PyTorch Precision Test')
    parser.add_argument('--cfg_file', type=str, default='/home/uestc/nvs/OpenStereo/cfgs/lightstereo/lightstereo_l_kitti.yaml', help='model config file path')
    parser.add_argument('--pth_weights', type=str, default='/home/uestc/nvs/OpenStereo/checkpoints/stereo_anything/StereoAnything-LightStereo_L.pt', help='PyTorch model weights path')
    parser.add_argument('--onnx_weights', type=str, default='/home/uestc/nvs/OpenStereo/checkpoints/stereo_anything/StereoAnything-LightStereo_L.onnx', help='ONNX model path')
    parser.add_argument('--left_img_path', type=str, default='/home/uestc/nvs/OpenStereo/datasets/Test_data/im0.png', help='left input image path')
    parser.add_argument('--right_img_path', type=str, default='/home/uestc/nvs/OpenStereo/datasets/Test_data/im1.png', help='right input image path')
    parser.add_argument('--output_dir', type=str, default='onnx_results', help='output directory for results')
    parser.add_argument('--num_tests', type=int, default=1, help='number of warmup runs for performance test')
    parser.add_argument('--vis_diff', action='store_true', help='visualize difference map')
    args = parser.parse_args()
    
    # Create output directory
    os.makedirs(args.output_dir, exist_ok=True)
    
    return args

def create_trainer(cfg_file, weights_path):
    """Create trainer and model using project's standard method"""
    # Create dummy args
    class Args:
        def __init__(self):
            self.cfg_file = cfg_file
            self.pretrained_model = weights_path
            self.dist_mode = False
            self.timeit = False
            self.run_mode = 'infer'
    
    args = Args()
    
    # Load config
    yaml_config = common_utils.config_loader(args.cfg_file)
    cfgs = EasyDict(yaml_config)
    
    # Set pretrained model path
    cfgs.MODEL.PRETRAINED_MODEL = weights_path
    
    # Create logger
    logger = common_utils.create_logger(log_file=None, rank=0)
    
    # Build trainer (this handles model creation and weight loading)
    trainer = build_trainer(args, cfgs, 0, 0, logger, None)
    
    return trainer.model

def load_data_transform(cfg_file_path):
    """Load data transform from config file"""
    yaml_config = common_utils.config_loader(cfg_file_path)
    cfgs = EasyDict(yaml_config)
    transform = build_transform_by_cfg(cfgs.DATA_CONFIG.DATA_TRANSFORM.EVALUATING)
    return transform

def load_input_images(args, transform):
    """Load and transform input images"""
    left_img = np.array(Image.open(args.left_img_path).convert('RGB'), dtype=np.float32)
    right_img = np.array(Image.open(args.right_img_path).convert('RGB'), dtype=np.float32)
    
    sample = transform({
        'left': left_img, 
        'right': right_img
    })
    
    # Add batch dimension
    return {
        'left': sample['left'].unsqueeze(0),
        'right': sample['right'].unsqueeze(0)
    }

def run_pytorch_inference(model, inputs):
    """Run PyTorch inference and return disparity map"""
    model.eval()
    with torch.no_grad(), torch.cuda.amp.autocast():
        inputs_gpu = {
            k: v.cuda() for k, v in inputs.items()
        }
        output = model(inputs_gpu)
    
    return output['disp_pred'].squeeze().cpu().numpy()

def run_onnx_inference(onnx_path, inputs):
    """Run ONNX inference and return disparity map"""
    # Prepare input data for ONNX
    ort_inputs = {
        'left_img': inputs['left'].numpy().astype(np.float32),
        'right_img': inputs['right'].numpy().astype(np.float32)
    }
    
    # Create ONNX Runtime session
    options = onnxruntime.SessionOptions()
    options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL
    providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
    
    session = onnxruntime.InferenceSession(onnx_path, options, providers)
    
    # Run inference
    ort_outputs = session.run(None, ort_inputs)

    print("ort_outputs_shape:",len(ort_outputs))

    # 打印输出形状
    print("\n[ONNX Output Shapes]")
    for i, output in enumerate(ort_outputs):
        print(f"Output {i}: {output.shape}")


    return ort_outputs[0].squeeze()



def calculate_metrics(pytorch_disp, onnx_disp):
    """Calculate metrics between PyTorch and ONNX outputs"""
    diff = np.abs(pytorch_disp - onnx_disp)
    metrics = {
        'EPE': np.mean(diff),
        'MaxDiff': np.max(diff),
        'MSE': np.mean(np.square(diff)),
        'Bad1': np.mean(diff > 1) * 100,
        'Bad3': np.mean(diff > 3) * 100
    }
    return metrics, diff

def benchmark_model(fn, *args, **kwargs):
    """Benchmark model inference time"""
    # Warmup runs
    for _ in range(5):
        fn(*args, **kwargs)
    
    # Timed runs
    times = []
    for _ in range(20):
        start_time = time.perf_counter()
        fn(*args, **kwargs)
        times.append(time.perf_counter() - start_time)
    
    return {
        'min': min(times) * 1000,
        'max': max(times) * 1000,
        'mean': sum(times) / len(times) * 1000
    }

def save_results(args, name, disparity, diff_map=None):
    """Save disparity and difference visualization"""
    # Save disparity as numpy array
    np.save(os.path.join(args.output_dir, f"{name}_disparity.npy"), disparity)
    
    # Save color visualization using plt.imsave
    plt.imsave(
        os.path.join(args.output_dir, f"{name}_color.png"),
        disparity,
        cmap='jet'
    )
    
    # Save difference map if provided
    if diff_map is not None:
        plt.imsave(
            os.path.join(args.output_dir, f"{name}_diff.png"), 
            np.clip(diff_map, 0, 10), 
            cmap='jet',
            vmin=0,
            vmax=10
        )
        np.save(os.path.join(args.output_dir, f"{name}_diff.npy"), diff_map)

def main():
    args = parse_config()
    
    # ==================== Prepare Data ====================
    print("Loading data transform...")
    transform = load_data_transform(args.cfg_file)
    inputs = load_input_images(args, transform)
    
    # ==================== PyTorch Model ====================
    print("\n[PyTorch Model Testing]")
    print(f"Loading PyTorch model from {args.pth_weights}...")
    pytorch_model = create_trainer(args.cfg_file, args.pth_weights)
    pytorch_model.cuda()
    
    # Run inference
    print("Running PyTorch inference...")
    pytorch_disp = run_pytorch_inference(pytorch_model, inputs)
    save_results(args, "pytorch", pytorch_disp)
    
    # Benchmark
    print("Benchmarking PyTorch model...")
    pytorch_stats = benchmark_model(
        run_pytorch_inference, 
        pytorch_model, 
        inputs
    )
    
    # ==================== ONNX Model ====================
    print("\n[ONNX Model Testing]")
    print(f"Loading ONNX model from {args.onnx_weights}...")
    
    # Run inference
    print("Running ONNX inference...")
    onnx_disp = run_onnx_inference(args.onnx_weights, inputs)
    save_results(args, "onnx", onnx_disp)
    
    # Benchmark
    print("Benchmarking ONNX model...")
    onnx_stats = benchmark_model(
        run_onnx_inference,
        args.onnx_weights,
        inputs
    )
    
    # ==================== Precision Comparison ====================
    print("\n[Precision Comparison]")
    
    # Calculate metrics
    metrics, diff_map = calculate_metrics(pytorch_disp, onnx_disp)
    
    # Save difference visualization
    if args.vis_diff:
        save_results(args, "diff", onnx_disp, diff_map)
    
    # ==================== Results Report ====================
    print("\n" + "="*60)
    print("Performance Comparison (ms per inference):")
    print("-"*60)
    print(f"{'Metric':<20}{'PyTorch':<15}{'ONNX':<15}{'Speedup':<15}")
    print(f"{'Min Latency':<20}{pytorch_stats['min']:.2f}{'ms':<10}{onnx_stats['min']:.2f}{'ms':<10}{(pytorch_stats['min']/onnx_stats['min']):.2f}x")
    print(f"{'Avg Latency':<20}{pytorch_stats['mean']:.2f}{'ms':<10}{onnx_stats['mean']:.2f}{'ms':<10}{(pytorch_stats['mean']/onnx_stats['mean']):.2f}x")
    print(f"{'Max Latency':<20}{pytorch_stats['max']:.2f}{'ms':<10}{onnx_stats['max']:.2f}{'ms':<10}{(pytorch_stats['max']/onnx_stats['max']):.2f}x")
    
    print("\nPrecision Metrics (PyTorch vs ONNX):")
    print("-"*60)
    for metric, value in metrics.items():
        print(f"{metric:<15}: {value:.6f}")
    
    print("\nOutput saved to:", args.output_dir)
    print("="*60)

if __name__ == '__main__':
    main()

RKNN推理结果问题 #15

Description

Metadata

Metadata

Assignees

Labels

Projects

Milestone

Relationships

Development

Issue actions