Kornia 如何用 PyTorch 搭起空间 AI 几何工具链

把 Kornia 理解成“移植到 PyTorch 的 OpenCV”，是最容易影响技术选型的误读。Kornia 的价值不只在于提供一批孤立的图像算子，更在于把几何处理变成 PyTorch 训练图中可以求导的一层：变换参数、单应矩阵、几何误差都能进入自动微分和优化器闭环。资料将其概括为基于 PyTorch 的可微几何计算机视觉库，并提供图像处理、几何变换、特征提取和模型训练等工具。

它与 OpenCV、PIL 更适合形成互补。OpenCV 和 PIL 仍适合解码、文件读写、显示以及那些不需要梯度的传统视觉流程；Kornia 更适合处理训练阶段的仿射、透视、增广、配准以及几何一致性损失。Kornia 不会让任意 OpenCV 操作自动获得梯度，能被优化的是 Kornia 自身实现的可微算子及其组成的计算图。 同时，Kornia 也不是开箱即用的完整 SLAM 或摄影测量系统，光束平差、三维重建及部分诊断工具仍需结合当前版本、具体模块和工程约束评估。

GitHub 仓库是核对版本、示例和 Issue 的实时入口，仓库对自身的定位是“Geometric Computer Vision Library for Spatial AI”。 资料中常见的“3000 星”来自 2020 年 11 月 3 日的报道，只能视为当时的热度快照，不能替代当前 GitHub 数据。 后续示例以一对图像 A、B 的几何对齐为主线，覆盖可微变换、配准、平面深度估计、特征审计和三维重建。

先让第一个可微几何变换跑起来

安装时先根据本机 CPU、CUDA 或其他加速环境选择合适的 PyTorch 版本，再安装 Kornia。不要一开始就把“最新版本”当作兼容版本，生产项目应记录实际解析到的包版本、Python 版本和设备环境。

python -m pip install torch
python -m pip install kornia
python -m pip freeze > requirements.lock.txt

安装后先检查实际导入的版本，而不是只看 pip list 中显示的需求约束：

python -c "from importlib.metadata import version; import torch, kornia; print('torch:', torch.__version__); print('kornia:', version('kornia')); print('cuda:', torch.cuda.is_available())"

Kornia 的 API 总体遵循 PyTorch 风格，支持自动微分和 GPU 执行，但具体参数、默认值和返回结构仍应以当前版本文档为准。 几何算子在不同版本间可能调整过边界模式名称、align_corners 默认值或归一化方式，因此直接复制旧代码前，应先检查函数签名：

python -c "import inspect, kornia as K; print(inspect.signature(K.geometry.transform.warp_affine)); print(inspect.signature(K.geometry.transform.warp_perspective))"

使用具体算子前，需先约定好张量规则。最常用的图像张量形状是 B×C×H×W，其中 B 是批次，C 是通道，H、W 是空间尺寸；几何矩阵通常也带批次维，例如仿射矩阵为 B×2×3、单应矩阵为 B×3×3。Kornia 的批处理能力继承自 PyTorch，这也是它适合在训练管线中批量执行几何操作的基础。

这段代码用于完成导入检查、单位变换、形状检查、梯度检查和有效区域掩码检查。为兼容常见的版本差异，_warp 会检测 border_mode 与 padding_mode 的命名：

import inspect
import torch
import kornia as K

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
torch.manual_seed(7)


def _warp(fn, src, matrix, dsize):
    parameters = inspect.signature(fn).parameters
    kwargs = {}

    if "align_corners" in parameters:
        kwargs["align_corners"] = True

    if "border_mode" in parameters:
        kwargs["border_mode"] = "zeros"
    elif "padding_mode" in parameters:
        kwargs["padding_mode"] = "zeros"

    return fn(src, matrix, dsize, **kwargs)


def warp_affine_compat(src, matrix, dsize):
    return _warp(
        K.geometry.transform.warp_affine,
        src,
        matrix,
        dsize,
    )


def warp_perspective_compat(src, matrix, dsize):
    return _warp(
        K.geometry.transform.warp_perspective,
        src,
        matrix,
        dsize,
    )


B, C, H, W = 2, 3, 64, 64
image = torch.rand(B, C, H, W, device=device)

identity = (
    torch.eye(3, device=device)[:2]
    .unsqueeze(0)
    .expand(B, -1, -1)
)
restored = warp_affine_compat(image, identity, (H, W))
assert restored.shape == image.shape
assert torch.allclose(restored, image, atol=1e-5)

matrix = torch.tensor(
    [[[1.0, 0.0, -4.0], [0.0, 1.0, 3.0]]],
    device=device,
).expand(B, -1, -1).clone().requires_grad_(True)

warped = warp_affine_compat(image, matrix, (H, W))
loss = warped.square().mean()
loss.backward()

assert warped.shape == image.shape
assert matrix.grad is not None
assert matrix.grad.shape == matrix.shape
assert torch.isfinite(matrix.grad).all()

valid = torch.ones(B, 1, H, W, device=device)
valid_warped = warp_affine_compat(valid, matrix, (H, W))

print("output:", tuple(warped.shape))
print("gradient:", tuple(matrix.grad.shape))
print("valid fraction:", valid_warped.mean().item())

这里的掩码没有假设某个版本的 warp_affine 一定接受 mask 参数，而是用同一矩阵变换一张全 1 图像，得到哪些采样点仍落在有效区域内。这种检查也能辅助诊断：如果掩码在平移后几乎全部归零，问题通常出在变换矩阵、边界模式或坐标约定上，而不是模型本身没有学起来。

模块选择应从具体需求出发，而不是按“功能数量”取舍：

如果需求只是“把图片旋转几度”或“按透视矩阵重采样”，应进入 kornia.geometry；如果需求是训练一个会同时改变图像和标签的增广策略，才有必要进一步看 kornia.augmentation。按需求拆分功能，比在一个庞杂的 API 列表中反复试错更有效。

把配准与深度估计接入训练闭环

推理时，可以把 Kornia 当作支持批处理的几何算子：输入图像和变换参数，输出变换后的张量。训练时则多了一层要求，几何参数必须保留梯度，计算出的几何误差要能通过 PyTorch 的自动微分回传，再用 Adam 等优化器更新。资料将自动微分、GPU/TPU 加速和批数据处理列为 Kornia 的主要特性，实际设备支持仍取决于 PyTorch 后端及具体算子。

可微数据增广是这种设计最直接的应用。资料特别提到，Kornia 不只是把传统算子换成 PyTorch 写法，还把裁切等操作表达为透视变换和仿射变换，使其可以参与基于梯度下降的优化。 早期资料将 TorchVision、Albumentations 视为不可微工具，并以此说明可微数据增广的动机；这反映的是当时的技术背景，不应被理解为对今天所有版本或第三方实现的永久判断。

双图配准示例不依赖外部图片，先构造一对有已知几何关系的合成图 A、B，再加入带噪声的对应点来初始化单应矩阵。实际项目只需将 A、B 替换为真实图像，并接入特征匹配结果。

A = torch.rand(1, 3, 96, 128, device=device)

src_pts = torch.tensor(
    [[[20.0, 20.0], [105.0, 18.0], [112.0, 78.0], [28.0, 86.0]]],
    device=device,
)
dst_pts = torch.tensor(
    [[[24.0, 18.0], [109.0, 20.0], [115.0, 80.0], [31.0, 88.0]]],
    device=device,
)

M_true = K.geometry.transform.get_perspective_transform(src_pts, dst_pts)
B = warp_perspective_compat(A, M_true, (96, 128))

noise = 0.03 * torch.randn_like(dst_pts)
M_init = K.geometry.transform.get_perspective_transform(
    src_pts,
    dst_pts + noise,
)

M = M_init.clone().requires_grad_(True)
optimizer = torch.optim.Adam([M], lr=1e-3)

mask = torch.zeros(1, 1, 96, 128, device=device)
mask[:, :, 10:-10, 10:-10] = 1.0

best_loss = torch.tensor(float("inf"), device=device)
M_refined = M.detach().clone()

for step in range(200):
    optimizer.zero_grad(set_to_none=True)

    warped = warp_perspective_compat(A, M, (96, 128))
    loss = (warped - B).abs().mean()

    loss.backward()

    assert M.grad is not None
    assert M.grad.shape == M.shape
    assert torch.isfinite(M.grad).all()

    if loss.detach() < best_loss:
        best_loss = loss.detach()
        M_refined = M.detach().clone()

    optimizer.step()

这一闭环依次包括四步：先用对应点或特征匹配得到单应矩阵初值，再用变换后的图像与目标图像计算光度或特征残差；随后通过 PyTorch 自动微分得到单应矩阵的梯度，最后更新参数，并检查残差、梯度以及有效区域。

这一环节最容易出错的是矩阵方向，而非优化器设置。get_perspective_transform、点变换和 warp_perspective 必须使用一致的像素坐标与归一化约定。不要仅凭“图像看起来稍微对齐”就继续调学习率，应先拿已知对应点做往返测试：如果点投影结果与目标点不一致，就要先检查矩阵方向、align_corners、原点位置和点坐标归一化方式。某些版本的点变换接口返回归一化坐标，另一些路径接收像素坐标，混用后会出现数量级看似不合理的损失。

单应矩阵还存在尺度不唯一性，比较矩阵元素绝对值通常没有意义。更有用的检查包括：对应点重投影误差、边界有效比例、矩阵条件数、光度残差，以及特征匹配中的内点比例。真实双图还应处理重复纹理、低纹理区域、动态物体和遮挡，不能把光度损失下降直接等同于几何关系正确。

深度估计需要先区分两件事：Kornia 提供的是可微几何工具，不是一个开箱即用的单目深度模型。预训练深度网络、相机标定和位姿仍然需要由外部提供；Kornia 可以介入数据增广、几何投影、深度变换和跨视图一致性约束。

一个可运行的深度基线是估计 fronto-parallel 平面的深度。设焦距为 (f)、双目基线为 (b)、平面深度为 (Z)，水平视差为 (fb/Z)。将这个视差写成可微平移矩阵，再用 Kornia 变换图像，就能让深度参数进入梯度闭环：

left = torch.rand(1, 3, 96, 128, device=device)
focal = 500.0
baseline = 0.1


def affine_shift(tx):
    linear = torch.tensor(
        [[[1.0, 0.0], [0.0, 1.0]]],
        device=left.device,
    )
    offset = torch.stack(
        (tx, torch.zeros_like(tx)),
        dim=0,
    ).reshape(1, 2, 1)
    return torch.cat((linear, offset), dim=-1)


z_true = torch.tensor(2.0, device=left.device)
tx_true = -focal * baseline / z_true
right = warp_affine_compat(
    left,
    affine_shift(tx_true),
    (96, 128),
)

depth_mask = torch.zeros(1, 1, 96, 128, device=left.device)
depth_mask[:, :, 20:-20, 30:-30] = 1.0

log_z = torch.log(torch.tensor(3.0, device=left.device))
optimizer = torch.optim.Adam([log_z], lr=5e-2)

for step in range(300):
    optimizer.zero_grad(set_to_none=True)

    z = log_z.exp()
    tx = -focal * baseline / z
    predicted = warp_affine_compat(
        left,
        affine_shift(tx),
        (96, 128),
    )

    loss = (predicted - right).abs().mean()
    loss.backward()
    optimizer.step()

z_estimate = log_z.detach().exp()

这段代码估计的是单个平面深度，而不是一般场景的深度图。它的作用是展示初始化深度、构造投影关系、计算图像残差和优化深度这条路径。对于真实任务，常见做法是让网络预测逆深度或视差，再结合相机模型将其转换为三维点或深度；Kornia 可以承担其中的可微投影、增广和一致性约束，但不会替代网络本身、标定文件或位姿估计。

推理阶段还可以关闭梯度以减少显存和计算开销：

with torch.no_grad():
    aligned = warp_perspective_compat(A, M_refined, (96, 128))

训练阶段则应保留几何矩阵的梯度，并避免在循环中用 .item() 读取损失，因为这类同步操作会削弱 GPU 流水线效率。

特征、光束平差、审计输出与三维重建

从双图配准走向三维重建，通常需要依次经过特征提取与描述、跨图匹配、几何验证、相对位姿估计、三角化、光束平差和场景重建。Kornia 的特征模块可以承担检测、描述和匹配。图像配准、深度估计、光束平差及场景重建构成一条由几何误差驱动的优化链，Kornia 的作用也不止于作为特征包。

这条路线中，特征匹配只是入口。重复纹理会产生错误对应点，低纹理区域则可能缺少足够约束，因此需要用单应矩阵、基本矩阵或本质矩阵等模型做几何验证，并区分内点与外点。进入三维阶段后，光束平差会联合调整相机参数和三维点，使重投影误差下降；场景重建则是这些几何结果的空间化表达。

这里需要区分“核心能力”和“实验性入口”。仿射、透视、滤波、颜色转换、特征处理和可微增广属于日常使用范围；光束平差、高级重建和诊断扩展是否位于 kornia.contrib，以及接口是否稳定，则应按当前版本文档核对。 参考资料将 Kornia 描述为覆盖从特征提取到场景重建的工具链，但实验模块的定位意味着它更适合作为路线图和原型工具，不能仅凭一篇教程就推断其已经满足生产级的大规模重建要求。

实验审计是 kornia.contrib 的一个典型用途。按资料提供的接口行为，summary() 会用一行概览捕获步骤数、几何状态和空间输入数量，同时明确给出 Image reconstruction not evaluated，随后附加告警；to_dict() 则可导出 JSON 安全字典，张量会附带 shape、dtype、device 元数据，非有限数值会被转换为 None。

审计接口可按以下方式使用；具体审计对象如何创建，应以当前 kornia.contrib 文档为准：

import json
from pathlib import Path

# audit 为按当前版本文档创建的 contrib 审计对象
print(audit.summary())

payload = audit.to_dict()
Path("geometry_audit.json").write_text(
    json.dumps(payload, ensure_ascii=False, indent=2, allow_nan=False),
    encoding="utf-8",
)

Image reconstruction not evaluated 意味着诊断对象只检查了被捕获的中间几何状态，并没有完成最终三维重建质量评估。它可以帮助发现坐标错位、状态缺失或非有限数值，但不能替代重投影误差、几何内点率、相机残差和三维点分布等最终指标。

工程实现还需处理性能、设备、复现性和传统库协作四类问题。批处理和 GPU 能提高吞吐，但高分辨率双图会同时放大特征、描述子和有效像素的内存占用，必要时应在检测、匹配和优化之间分段处理。几何矩阵、图像和优化器状态必须保持同一设备，只有解码、显示和文件导出才需要回到 CPU。

复现性方面，除了固定随机种子，还应记录 PyTorch 与 Kornia 的准确版本、CUDA 环境、图像缩放方式、坐标约定、边界模式和增广参数。与 OpenCV、PIL 协作时，较合理的分工是由 OpenCV 或 PIL 负责解码和显示，由 Kornia 负责进入训练图的几何处理；如果必须在中途调用传统视觉算法，应明确该步骤已经切断梯度图，不能再期待几何误差穿过它反向传播。

完成配准后，运行一次 A、B 双图配准，保存叠加图、绝对残差和包含单应矩阵的 JSON 几何状态文件：

import json
import math
import matplotlib.pyplot as plt

aligned = warp_perspective_compat(
    A,
    M_refined,
    (A.shape[-2], A.shape[-1]),
)

residual = (aligned - B).abs()

final_mask = torch.zeros(
    1,
    1,
    A.shape[-2],
    A.shape[-1],
    device=A.device,
)
final_mask[:, :, 10:-10, 10:-10] = 1.0
valid = final_mask.expand_as(residual)
residual_mae = (residual * valid).sum() / valid.sum()

overlay = 0.5 * (aligned + B)
overlay_vis = (
    overlay[0]
    .permute(1, 2, 0)
    .clamp(0.0, 1.0)
    .cpu()
    .numpy()
)
residual_vis = (
    residual.mean(dim=1)[0]
    .cpu()
    .numpy()
)

fig, axes = plt.subplots(1, 2, figsize=(10, 4))
axes[0].imshow(overlay_vis)
axes[0].set_title("Aligned overlay")
axes[1].imshow(residual_vis, cmap="magma")
axes[1].set_title("Absolute residual")
fig.tight_layout()
fig.savefig(
    "registration_overlay_residual.png",
    dpi=150,
    bbox_inches="tight",
)
plt.close(fig)


def finite_or_none(value):
    value = float(value)
    return value if math.isfinite(value) else None


matrix_json = [
    [finite_or_none(value) for value in row]
    for row in M_refined[0].detach().cpu().tolist()
]

geometry_state = {
    "homography": matrix_json,
    "residual_mae": finite_or_none(residual_mae),
    "valid_pixels": int(valid.sum().item()),
    "input_shape": list(A.shape),
    "dtype": str(A.dtype),
    "device": str(A.device),
}

with open("registration_state.json", "w", encoding="utf-8") as file:
    json.dump(
        geometry_state,
        file,
        ensure_ascii=False,
        indent=2,
        allow_nan=False,
    )

print("saved registration_overlay_residual.png")
print("saved registration_state.json")