GSplat项目中张量尺寸不匹配问题的分析与解决

问题背景

在GSplat项目的测试过程中，发现test_fully_fused_projection_packed测试用例存在间歇性失败的问题。该测试用于验证投影变换的核心功能，涉及多个张量操作和梯度计算。失败时会出现张量尺寸不匹配的错误，具体表现为两个张量在非单一维度上的尺寸不一致。

问题现象

测试失败时抛出的错误信息显示，在计算梯度时，张量a和张量b在第0维的尺寸不一致。例如，一个张量尺寸为193831，而另一个为193829。这种尺寸差异导致无法完成张量间的乘法操作。

技术分析

该测试用例的核心是验证投影变换及其梯度计算的正确性。测试过程中会生成随机数据，包括视图矩阵(viewmats)、四元数(quats)、缩放因子(scales)和均值(means)，然后通过这些参数计算2D均值(means2d)、深度(depths)和圆锥曲线(conics)。

问题出现在梯度计算阶段，当使用选择索引sel来筛选有效数据时，可能导致前后张量尺寸不一致。这是因为sel是基于半径筛选的条件，而不同运行环境下随机生成的数据可能导致筛选结果存在微小差异。

解决方案

经过分析，提出了以下解决方案：

1. 使用__radii > 0作为筛选条件替代原来的sel索引。因为__radii直接反映了高斯分布的有效性，可以确保前后张量尺寸的一致性。

2. 修改梯度计算部分的代码，确保所有参与计算的张量在相同条件下进行筛选：

v_viewmats, v_quats, v_scales, v_means = torch.autograd.grad(
    (means2d * v_means2d[__radii > 0]).sum()
    + (depths * v_depths[__radii > 0]).sum()
    + (normals * v_normals[__radii > 0]).sum()
    + (conics * v_conics[__radii > 0]).sum(),
    (viewmats, quats, scales, means),
    retain_graph=True,
)

技术原理

这个问题的本质在于PyTorch张量操作的尺寸一致性要求。在自动微分过程中，所有参与计算的张量必须在非单一维度上保持相同尺寸。使用基于半径的筛选条件比使用中间变量索引更可靠，因为它直接反映了数据的有效性边界。

影响范围

该问题主要影响测试环节，不会影响核心功能的正确性。但由于测试是保证代码质量的重要手段，修复这个问题对于确保项目稳定性非常重要。

总结

通过分析GSplat项目中出现的张量尺寸不匹配问题，我们理解了在复杂张量操作中保持尺寸一致性的重要性。使用更直接的筛选条件可以避免因中间变量导致的尺寸不一致问题。这个案例也提醒我们，在编写涉及多步骤张量操作的代码时，需要特别注意各步骤间的尺寸匹配问题。