test_precision.py

#!/usr/bin/env python3
"""
Grid Sampler CUDA 精度测试用例
测试数值精度、边界条件、极端值等场景
"""

import numpy as np
import sys
import os
import math
import time

# 添加当前目录到Python路径
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

try:
    import grid_sampler_cuda
except ImportError as e:
    print(f"无法导入grid_sampler_cuda模块: {e}")
    print("请先编译CUDA扩展")
    sys.exit(1)

# 尝试导入PyTorch作为参考实现
try:
    import torch
    import torch.nn.functional as F
    TORCH_AVAILABLE = True
except ImportError:
    TORCH_AVAILABLE = False
    print("警告: PyTorch未安装，将跳过与参考实现的对比测试")

def test_bilinear_precision():
    """测试双线性插值的数值精度"""
    print("=== 测试双线性插值数值精度 ===")

    # 创建已知的测试数据，便于验证精度
    input_data = np.array([[[
        [1.0, 2.0, 3.0],
        [4.0, 5.0, 6.0],
        [7.0, 8.0, 9.0]
    ]]], dtype=np.float32)  # [1, 1, 3, 3]

    # 测试精确的网格点（应该返回精确值）
    test_cases = [
        # (grid_x, grid_y, expected_value, description)
        (-1.0, -1.0, 1.0, "左上角"),
        (1.0, -1.0, 3.0, "右上角"),
        (-1.0, 1.0, 7.0, "左下角"),
        (1.0, 1.0, 9.0, "右下角"),
        (0.0, 0.0, 5.0, "中心点"),
    ]

    for grid_x, grid_y, expected, desc in test_cases:
        grid = np.array([[[[grid_x, grid_y]]]], dtype=np.float32)
        output = grid_sampler_cuda.grid_sampler(
            input_data, grid, mode="bilinear", align_corners=True
        )
        actual = output[0, 0, 0, 0]
        error = abs(actual - expected)

        print(f"{desc}: 期望={expected:.6f}, 实际={actual:.6f}, 误差={error:.2e}")

        # 对于精确的网格点，误差应该非常小
        if error > 1e-6:
            print(f"❌ 精度测试失败: {desc}")
            return False

    # 测试插值精度
    print("\n测试插值精度:")
    # 在(0,0)和(1,1)之间插值，期望得到5.0
    grid = np.array([[[[0.0, 0.0]]]], dtype=np.float32)
    output = grid_sampler_cuda.grid_sampler(
        input_data, grid, mode="bilinear", align_corners=True
    )
    actual = output[0, 0, 0, 0]
    expected = 5.0  # 中心值
    error = abs(actual - expected)
    print(f"中心插值: 期望={expected:.6f}, 实际={actual:.6f}, 误差={error:.2e}")

    if error > 1e-6:
        print("❌ 插值精度测试失败")
        return False

    print("✓ 双线性插值精度测试通过\n")
    return True

def test_boundary_precision():
    """测试边界条件的精度"""
    print("=== 测试边界条件精度 ===")

    # 创建简单的测试数据
    input_data = np.array([[[
        [1.0, 2.0],
        [3.0, 4.0]
    ]]], dtype=np.float32)  # [1, 1, 2, 2]

    # 测试边界填充模式
    padding_modes = ["zeros", "border", "reflection"]

    for padding_mode in padding_modes:
        print(f"\n测试 {padding_mode} 填充模式:")

        # 测试边界外的点
        test_cases = [
            (-2.0, -2.0, "左上角外"),
            (2.0, -2.0, "右上角外"),
            (-2.0, 2.0, "左下角外"),
            (2.0, 2.0, "右下角外"),
        ]

        for grid_x, grid_y, desc in test_cases:
            grid = np.array([[[[grid_x, grid_y]]]], dtype=np.float32)
            output = grid_sampler_cuda.grid_sampler(
                input_data, grid, mode="bilinear",
                padding_mode=padding_mode, align_corners=True
            )
            actual = output[0, 0, 0, 0]

            # 根据填充模式计算期望值
            if padding_mode == "zeros":
                expected = 0.0
            elif padding_mode == "border":
                # 边界填充应该返回最近的边界值
                if grid_x < -1: grid_x = -1
                if grid_x > 1: grid_x = 1
                if grid_y < -1: grid_y = -1
                if grid_y > 1: grid_y = 1
                # 简化的期望值计算
                expected = 2.5  # 近似值
            elif padding_mode == "reflection":
                # 反射填充
                expected = 2.5  # 近似值

            error = abs(actual - expected)
            print(f"  {desc}: 实际={actual:.6f}, 期望≈{expected:.6f}, 误差={error:.2e}")

    print("✓ 边界条件精度测试通过\n")
    return True

def test_align_corners_precision():
    """测试align_corners参数的精度影响"""
    print("=== 测试align_corners精度 ===")

    # 创建测试数据
    input_data = np.array([[[
        [1.0, 2.0, 3.0],
        [4.0, 5.0, 6.0],
        [7.0, 8.0, 9.0]
    ]]], dtype=np.float32)  # [1, 1, 3, 3]

    # 测试相同的网格坐标在不同align_corners设置下的结果
    test_grids = [
        (-1.0, -1.0, "左上角"),
        (0.0, 0.0, "中心"),
        (1.0, 1.0, "右下角"),
    ]

    for grid_x, grid_y, desc in test_grids:
        grid = np.array([[[[grid_x, grid_y]]]], dtype=np.float32)

        # align_corners=False
        output_false = grid_sampler_cuda.grid_sampler(
            input_data, grid, mode="bilinear", align_corners=False
        )
        result_false = output_false[0, 0, 0, 0]

        # align_corners=True
        output_true = grid_sampler_cuda.grid_sampler(
            input_data, grid, mode="bilinear", align_corners=True
        )
        result_true = output_true[0, 0, 0, 0]

        print(f"{desc}: align_corners=False={result_false:.6f}, align_corners=True={result_true:.6f}")

        # 验证两种模式产生不同的结果（除了特殊情况）
        if abs(result_false - result_true) < 1e-6:
            print(f"  注意: 两种模式结果相同（可能是特殊情况）")

    print("✓ align_corners精度测试通过\n")
    return True

def test_extreme_values():
    """测试极端值的处理"""
    print("=== 测试极端值处理 ===")

    # 测试极小值
    print("测试极小值:")
    input_data = np.array([[[[1e-10, 2e-10], [3e-10, 4e-10]]]], dtype=np.float32)
    grid = np.array([[[[-1.0, -1.0]]]], dtype=np.float32)
    output = grid_sampler_cuda.grid_sampler(input_data, grid, mode="bilinear")
    result = output[0, 0, 0, 0]
    print(f"  极小值输入: {result:.2e}")

    # 测试极大值
    print("测试极大值:")
    input_data = np.array([[[[1e10, 2e10], [3e10, 4e10]]]], dtype=np.float32)
    grid = np.array([[[[-1.0, -1.0]]]], dtype=np.float32)
    output = grid_sampler_cuda.grid_sampler(input_data, grid, mode="bilinear")
    result = output[0, 0, 0, 0]
    print(f"  极大值输入: {result:.2e}")

    # 测试零值
    print("测试零值:")
    input_data = np.zeros((1, 1, 2, 2), dtype=np.float32)
    grid = np.array([[[[-1.0, -1.0]]]], dtype=np.float32)
    output = grid_sampler_cuda.grid_sampler(input_data, grid, mode="bilinear")
    result = output[0, 0, 0, 0]
    print(f"  零值输入: {result:.2e}")

    # 测试负值
    print("测试负值:")
    input_data = np.array([[[[-1.0, -2.0], [-3.0, -4.0]]]], dtype=np.float32)
    grid = np.array([[[[-1.0, -1.0]]]], dtype=np.float32)
    output = grid_sampler_cuda.grid_sampler(input_data, grid, mode="bilinear")
    result = output[0, 0, 0, 0]
    print(f"  负值输入: {result:.2e}")

    print("✓ 极端值测试通过\n")
    return True

def test_accumulated_error():
    """测试累积误差"""
    print("=== 测试累积误差 ===")

    # 创建初始数据
    input_data = np.array([[[
        [1.0, 2.0, 3.0, 4.0],
        [5.0, 6.0, 7.0, 8.0],
        [9.0, 10.0, 11.0, 12.0],
        [13.0, 14.0, 15.0, 16.0]
    ]]], dtype=np.float32)  # [1, 1, 4, 4]

    # 进行多次变换，每次都是恒等变换
    current_data = input_data.copy()
    num_iterations = 10

    # 创建恒等变换网格
    grid = np.array([[[
        [[-1, -1], [-1, 1]],
        [[1, -1], [1, 1]]
    ]]], dtype=np.float32)  # 2x2输出

    print(f"进行{num_iterations}次恒等变换...")

    for i in range(num_iterations):
        # 将输出重新采样回原始大小
        # 这里简化处理，直接使用原始网格
        current_data = grid_sampler_cuda.grid_sampler(
            current_data, grid, mode="bilinear", align_corners=True
        )

        # 计算与原始数据的误差
        if i % 2 == 0:  # 每两次检查一次
            # 重新采样回原始大小进行比较
            original_grid = np.array([[[
                [[-1, -1], [-1, 1], [1, -1], [1, 1]],
                [[-1, -1], [-1, 1], [1, -1], [1, 1]],
                [[-1, -1], [-1, 1], [1, -1], [1, 1]],
                [[-1, -1], [-1, 1], [1, -1], [1, 1]]
            ]]], dtype=np.float32)

            resampled = grid_sampler_cuda.grid_sampler(
                current_data, original_grid, mode="bilinear", align_corners=True
            )

            error = np.mean(np.abs(resampled - input_data))
            print(f"  第{i+1}次变换后误差: {error:.2e}")

    print("✓ 累积误差测试通过\n")
    return True

def test_pytorch_comparison():
    """与PyTorch参考实现对比"""
    if not TORCH_AVAILABLE:
        print("=== 跳过PyTorch对比测试（PyTorch未安装） ===\n")
        return True

    print("=== 与PyTorch参考实现对比 ===")

    # 创建测试数据
    input_data = np.random.randn(1, 3, 8, 8).astype(np.float32)
    grid_data = np.random.randn(1, 4, 4, 2).astype(np.float32)
    grid_data = np.clip(grid_data, -1, 1)  # 限制在[-1, 1]范围内

    # CUDA实现
    cuda_output = grid_sampler_cuda.grid_sampler(
        input_data, grid_data, mode="bilinear", padding_mode="zeros", align_corners=False
    )

    # PyTorch实现
    torch_input = torch.from_numpy(input_data)
    torch_grid = torch.from_numpy(grid_data)
    torch_output = F.grid_sample(
        torch_input, torch_grid, mode="bilinear", padding_mode="zeros", align_corners=False
    )
    torch_output = torch_output.numpy()

    # 计算误差
    error = np.mean(np.abs(cuda_output - torch_output))
    max_error = np.max(np.abs(cuda_output - torch_output))

    print(f"与PyTorch对比:")
    print(f"  平均误差: {error:.2e}")
    print(f"  最大误差: {max_error:.2e}")
    print(f"  相对误差: {error / np.mean(np.abs(torch_output)):.2e}")

    # 检查误差是否在可接受范围内
    if error > 1e-5:
        print(f"❌ 与PyTorch对比误差过大: {error:.2e}")
        return False

    print("✓ PyTorch对比测试通过\n")
    return True

def test_numerical_stability():
    """测试数值稳定性"""
    print("=== 测试数值稳定性 ===")

    # 测试不同数据范围的稳定性
    test_ranges = [
        (1e-6, 1e-3, "极小值范围"),
        (1e-3, 1.0, "小值范围"),
        (1.0, 1e3, "中等值范围"),
        (1e3, 1e6, "大值范围"),
    ]

    for min_val, max_val, desc in test_ranges:
        print(f"测试{desc}: [{min_val:.0e}, {max_val:.0e}]")

        # 创建测试数据
        input_data = np.random.uniform(min_val, max_val, (1, 1, 4, 4)).astype(np.float32)
        grid_data = np.random.uniform(-1, 1, (1, 2, 2, 2)).astype(np.float32)

        # 多次运行，检查结果一致性
        results = []
        for _ in range(5):
            output = grid_sampler_cuda.grid_sampler(
                input_data, grid_data, mode="bilinear"
            )
            results.append(output.copy())

        # 检查结果一致性
        max_variation = 0
        for i in range(1, len(results)):
            variation = np.max(np.abs(results[i] - results[0]))
            max_variation = max(max_variation, variation)

        print(f"  最大变化: {max_variation:.2e}")

        if max_variation > 1e-6:
            print(f"❌ 数值稳定性测试失败: {desc}")
            return False

    print("✓ 数值稳定性测试通过\n")
    return True

def test_precision_consistency():
    """测试精度一致性"""
    print("=== 测试精度一致性 ===")

    # 创建测试数据
    input_data = np.array([[[
        [1.0, 2.0, 3.0],
        [4.0, 5.0, 6.0],
        [7.0, 8.0, 9.0]
    ]]], dtype=np.float32)

    # 测试不同模式下的精度一致性
    modes = ["bilinear", "nearest"]
    padding_modes = ["zeros", "border", "reflection"]

    # 创建测试网格
    grid_data = np.array([[[
        [[-0.5, -0.5], [0.5, -0.5]],
        [[-0.5, 0.5], [0.5, 0.5]]
    ]]], dtype=np.float32)

    results = {}

    for mode in modes:
        for padding_mode in padding_modes:
            output = grid_sampler_cuda.grid_sampler(
                input_data, grid_data, mode=mode, padding_mode=padding_mode
            )
            results[(mode, padding_mode)] = output.copy()

    # 检查结果的一致性（相同输入应该产生相同输出）
    print("检查结果一致性:")
    for key, result in results.items():
        mode, padding_mode = key
        # 多次运行相同输入
        for _ in range(3):
            new_result = grid_sampler_cuda.grid_sampler(
                input_data, grid_data, mode=mode, padding_mode=padding_mode
            )
            diff = np.max(np.abs(new_result - result))
            if diff > 1e-7:
                print(f"❌ 一致性测试失败: {mode}+{padding_mode}, 差异={diff:.2e}")
                return False

    print("✓ 精度一致性测试通过\n")
    return True

def main():
    """运行所有精度测试"""
    print("开始运行Grid Sampler CUDA精度测试...\n")

    test_functions = [
        test_bilinear_precision,
        test_boundary_precision,
        test_align_corners_precision,
        test_extreme_values,
        test_accumulated_error,
        test_pytorch_comparison,
        test_numerical_stability,
        test_precision_consistency,
    ]

    passed = 0
    total = len(test_functions)

    for test_func in test_functions:
        try:
            if test_func():
                passed += 1
            else:
                print(f"❌ {test_func.__name__} 失败")
        except Exception as e:
            print(f"❌ {test_func.__name__} 异常: {e}")
            import traceback
            traceback.print_exc()

    print(f"精度测试结果: {passed}/{total} 通过")

    if passed == total:
        print("🎉 所有精度测试通过！")
        return 0
    else:
        print(f"❌ {total - passed} 个测试失败")
        return 1

if __name__ == "__main__":
    exit(main())