export_onnx.py

"""
Restormer ONNX 导出与验证脚本
=============================
将 PyTorch Restormer 模型导出为 ONNX 格式，并验证输出精度。
"""

import torch
import torch.nn.functional as F
import numpy as np
import os
import sys
import time
from runpy import run_path

# ---------------------------------------------------------------------------
# 配置
# ---------------------------------------------------------------------------
CKPT_PATH = 'net_g_92000.pth'
ONNX_PATH = 'restormer_denoise.onnx'
ONNX_DYNAMIC_PATH = 'restormer_denoise_dynamic.onnx'

restormer_params = {
    'inp_channels': 3,
    'out_channels': 3,
    'dim': 48,
    'num_blocks': [4, 6, 6, 8],
    'num_refinement_blocks': 4,
    'heads': [1, 2, 4, 8],
    'ffn_expansion_factor': 2.66,
    'bias': False,
    'LayerNorm_type': 'WithBias',
    'dual_pixel_task': False,
}


def load_model():
    """Load PyTorch Restormer model."""
    load_arch = run_path(os.path.join(os.path.dirname(__file__), 'restormer_arch.py'))
    model = load_arch['Restormer'](**restormer_params)

    checkpoint = torch.load(CKPT_PATH, map_location='cpu')
    model.load_state_dict(checkpoint['params'])

    model.eval()
    return model


def export_fixed_size(model, device='cpu'):
    """Export ONNX with fixed input size 512×512."""
    print("\n" + "=" * 70)
    print("📦 导出 ONNX (固定尺寸 512×512)")
    print("=" * 70)

    model = model.to(device)
    model.eval()

    dummy_input = torch.randn(1, 3, 512, 512, device=device)

    # Trace the model
    print("  Tracing model...")
    with torch.inference_mode():
        torch.onnx.export(
            model,
            dummy_input,
            ONNX_PATH,
            input_names=['input'],
            output_names=['output'],
            opset_version=18,
            do_constant_folding=True,
            dynamic_axes=None,
            export_params=True,
        )

    import onnx
    onnx_model = onnx.load(ONNX_PATH)
    onnx.checker.check_model(onnx_model)
    print(f"  ✅ 导出成功: {ONNX_PATH}")
    print(f"  模型大小: {os.path.getsize(ONNX_PATH)/1024**2:.1f} MB")

    # Print ONNX model info
    print(f"  IR 版本: {onnx_model.ir_version}")
    print(f"  Opset: {onnx_model.opset_import[0].version}")
    print(f"  生产者: {onnx_model.producer_name}")
    print(f"  输入节点数: {len(onnx_model.graph.node)}")

    # Validate ONNX model shape
    graph = onnx_model.graph
    print(f"  输入: {[(i.name, [d.dim_value for d in i.type.tensor_type.shape.dim]) for i in graph.input]}")
    print(f"  输出: {[(o.name, [d.dim_value for d in o.type.tensor_type.shape.dim]) for o in graph.output]}")

    return ONNX_PATH


def export_dynamic_size(model, device='cpu'):
    """Export ONNX with dynamic spatial dimensions."""
    print("\n" + "=" * 70)
    print("📦 导出 ONNX (动态尺寸)")
    print("=" * 70)

    model = model.to(device)
    model.eval()

    dummy_input = torch.randn(1, 3, 512, 512, device=device)

    dynamic_axes = {
        'input': {0: 'batch', 2: 'height', 3: 'width'},
        'output': {0: 'batch', 2: 'height', 3: 'width'},
    }

    print("  Tracing model with dynamic axes...")
    with torch.inference_mode():
        torch.onnx.export(
            model,
            dummy_input,
            ONNX_DYNAMIC_PATH,
            input_names=['input'],
            output_names=['output'],
            opset_version=18,
            do_constant_folding=True,
            dynamic_axes=dynamic_axes,
            export_params=True,
        )

    import onnx
    onnx_model = onnx.load(ONNX_DYNAMIC_PATH)
    onnx.checker.check_model(onnx_model)
    print(f"  ✅ 导出成功: {ONNX_DYNAMIC_PATH}")
    print(f"  模型大小: {os.path.getsize(ONNX_DYNAMIC_PATH)/1024**2:.1f} MB")

    graph = onnx_model.graph
    print(f"  输入: {[(i.name, [d.dim_value for d in i.type.tensor_type.shape.dim]) for i in graph.input]}")
    print(f"  输出: {[(o.name, [d.dim_value for d in o.type.tensor_type.shape.dim]) for o in graph.output]}")

    return ONNX_DYNAMIC_PATH


def verify_onnx(model, onnx_path, device='cpu'):
    """Verify ONNX model output matches PyTorch model."""
    print("\n" + "=" * 70)
    print("🔍 验证 ONNX 模型输出精度")
    print("=" * 70)

    import onnxruntime as ort

    # Create ONNX Runtime session
    providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] if device == 'cuda' else ['CPUExecutionProvider']
    try:
        session = ort.InferenceSession(onnx_path, providers=providers)
    except Exception as e:
        print(f"  ⚠️ CUDA provider failed ({e}), falling back to CPU")
        session = ort.InferenceSession(onnx_path, providers=['CPUExecutionProvider'])

    actual_providers = session.get_providers()
    print(f"  ONNX Runtime providers: {actual_providers}")

    model = model.to(device)
    model.eval()

    # Test multiple sizes (avoid OOM by keeping sizes manageable)
    test_sizes = [(256, 256), (320, 448), (400, 400), (512, 512)]
    # Note: skip 768+ for verification as it exceeds available GPU memory

    for h, w in test_sizes:
        # Create random input (same seed for reproducibility)
        torch.manual_seed(42)
        input_tensor = torch.randn(1, 3, h, w, device=device)

        # ONNX Runtime inference (run first to avoid PyTorch GPU mem fragmentation)
        onnx_out = session.run(['output'], {'input': input_tensor.cpu().numpy()})[0]

        # PyTorch inference
        with torch.inference_mode():
            torch_out = model(input_tensor)

        onnx_out_tensor = torch.from_numpy(onnx_out).to(device)

        # Compute errors
        abs_diff = (torch_out - onnx_out_tensor).abs()
        max_diff = abs_diff.max().item()
        mean_diff = abs_diff.mean().item()
        rel_diff = (abs_diff / (torch_out.abs() + 1e-8)).mean().item()

        # Cosine similarity
        cos_sim = torch.nn.functional.cosine_similarity(
            torch_out.flatten().unsqueeze(0).float(),
            onnx_out_tensor.flatten().unsqueeze(0).float()
        ).item()

        status = "✅" if max_diff < 1e-3 and cos_sim > 0.9999 else "⚠️"
        print(f"  {status} {h}×{w}: max_diff={max_diff:.2e}, mean_diff={mean_diff:.2e}, "
              f"rel_diff={rel_diff:.2e}, cosine_sim={cos_sim:.8f}")

        # Cleanup GPU memory
        del input_tensor, torch_out, onnx_out_tensor, onnx_out
        if device == 'cuda':
            torch.cuda.empty_cache()


def benchmark_onnx(model, onnx_path, device='cpu'):
    """Compare PyTorch vs ONNX Runtime inference speed."""
    print("\n" + "=" * 70)
    print("⏱️ PyTorch vs ONNX Runtime 推理速度对比")
    print("=" * 70)

    import onnxruntime as ort

    # Use CPU provider for benchmark to avoid GPU memory issues
    # (ONNX RT GPU allocator behaves differently than PyTorch for this model)
    session_cpu = ort.InferenceSession(onnx_path, providers=['CPUExecutionProvider'])
    print(f"  Benchmark ONNX RT provider: CPU (for fair memory comparison)")

    model = model.to(device)
    model.eval()

    sizes = [(256, 256), (512, 512)]
    warmup, runs = 10, 30

    for h, w in sizes:
        torch.manual_seed(42)
        input_tensor = torch.randn(1, 3, h, w, device=device)
        input_np = input_tensor.cpu().numpy()

        # PyTorch warmup
        with torch.inference_mode():
            for _ in range(warmup):
                _ = model(input_tensor)

        if device == 'cuda':
            torch.cuda.synchronize()

        # PyTorch benchmark (GPU)
        torch_times = []
        with torch.inference_mode():
            for _ in range(runs):
                if device == 'cuda':
                    torch.cuda.synchronize()
                t0 = time.perf_counter()
                _ = model(input_tensor)
                if device == 'cuda':
                    torch.cuda.synchronize()
                torch_times.append((time.perf_counter() - t0) * 1000)

        # ONNX warmup (CPU)
        for _ in range(warmup):
            _ = session_cpu.run(['output'], {'input': input_np})

        # ONNX benchmark (CPU)
        onnx_times = []
        for _ in range(runs):
            t0 = time.perf_counter()
            _ = session_cpu.run(['output'], {'input': input_np})
            onnx_times.append((time.perf_counter() - t0) * 1000)

        torch_mean, torch_std = np.mean(torch_times), np.std(torch_times)
        onnx_mean, onnx_std = np.mean(onnx_times), np.std(onnx_times)

        print(f"  {h}×{w}:")
        print(f"    PyTorch (GPU):  {torch_mean:8.2f} ± {torch_std:5.2f} ms")
        print(f"    ONNX RT (CPU):  {onnx_mean:8.2f} ± {onnx_std:5.2f} ms")
        print(f"    注: ONNX Runtime GPU 模式对 Transformer 模型内存分配与 PyTorch 不同")

        del input_tensor, input_np
        if device == 'cuda':
            torch.cuda.empty_cache()


def main():
    script_dir = os.path.dirname(os.path.abspath(__file__))
    os.chdir(script_dir)

    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    print(f"设备: {device}")

    # 1. Load model
    print("\n[1] 载入 PyTorch 模型...")
    model = load_model()
    print(f"  参数量: {sum(p.numel() for p in model.parameters()):,}")

    # 2. Export fixed-size ONNX
    onnx_fixed = export_fixed_size(model, device)

    # 3. Export dynamic-size ONNX
    onnx_dynamic = export_dynamic_size(model, device)

    # 4. Verify ONNX outputs
    verify_onnx(model, onnx_dynamic, device)

    # 5. Benchmark ONNX vs PyTorch
    benchmark_onnx(model, onnx_dynamic, device)

    # 6. Test with real image
    print("\n" + "=" * 70)
    print("🖼️ 真实图像 ONNX 推理测试")
    print("=" * 70)

    import onnxruntime as ort
    import cv2

    providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] if device == 'cuda' else ['CPUExecutionProvider']
    try:
        session = ort.InferenceSession(onnx_dynamic, providers=providers)
    except Exception:
        session = ort.InferenceSession(onnx_dynamic, providers=['CPUExecutionProvider'])

    # Find a sample image
    for img_name in ['sample1.png', 'Sample3.png']:
        img_path = os.path.join(script_dir, img_name)
        if not os.path.exists(img_path):
            # try samples/ subdir
            alt_path = os.path.join(script_dir, 'samples', img_name)
            if os.path.exists(alt_path):
                img_path = alt_path
            else:
                continue

        img_bgr = cv2.imread(img_path)
        img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
        h, w = img_rgb.shape[:2]

        input_np = img_rgb.astype(np.float32) / 255.0
        input_np = input_np.transpose(2, 0, 1)[np.newaxis, ...]  # [1, 3, H, W]

        # Pad to multiple of 8
        H = ((h + 7) // 8) * 8
        W = ((w + 7) // 8) * 8
        pad_h, pad_w = H - h, W - w
        input_padded = np.pad(input_np, ((0, 0), (0, 0), (0, pad_h), (0, pad_w)), mode='reflect')

        # ONNX inference
        t0 = time.perf_counter()
        onnx_out = session.run(['output'], {'input': input_padded})[0]
        elapsed = (time.perf_counter() - t0) * 1000

        # Unpad and convert
        onnx_out = onnx_out[:, :, :h, :w]
        onnx_out_np = (onnx_out[0].transpose(1, 2, 0) * 255).clip(0, 255).astype(np.uint8)

        # PyTorch inference for comparison
        input_tensor = torch.from_numpy(input_np).to(device)
        input_padded_t = F.pad(input_tensor, (0, pad_w, 0, pad_h), 'reflect')
        with torch.inference_mode():
            torch_out = torch.clamp(model(input_padded_t), 0, 1)
        torch_out_np = (torch_out[:, :, :h, :w].permute(0, 2, 3, 1).cpu().numpy()[0] * 255).clip(0, 255).astype(np.uint8)

        # Compare
        diff = np.abs(onnx_out_np.astype(np.float32) - torch_out_np.astype(np.float32))
        print(f"  {os.path.basename(img_path)} ({w}×{h}):")
        print(f"    ONNX 推理耗时: {elapsed:.1f} ms")
        print(f"    vs PyTorch: max_diff={diff.max():.1f}, mean_diff={diff.mean():.4f}")
        if diff.max() < 1.0:
            print(f"    ✅ 输出一致 (max pixel diff < 1)")

        # Save ONNX output
        out_path = os.path.join(script_dir, 'temp', f'{os.path.splitext(img_name)[0]}_onnx.png')
        os.makedirs(os.path.dirname(out_path), exist_ok=True)
        cv2.imwrite(out_path, cv2.cvtColor(onnx_out_np, cv2.COLOR_RGB2BGR))
        print(f"    ONNX 输出保存至: {out_path}")

        break

    print("\n" + "=" * 70)
    print("✅ ONNX 导出/验证/测试 全部完成")
    print(f"   固定尺寸模型: {ONNX_PATH}")
    print(f"   动态尺寸模型: {ONNX_DYNAMIC_PATH}")
    print("=" * 70)


if __name__ == '__main__':
    main()