PyTorch Spatial Transformer Network输出张量全零问题排查

核心排查方向及修复步骤 #

• 仿射变换参数异常
  STN的定位网络（Localization Network）输出的theta矩阵如果出现极端值，会让所有有效像素都被移出输出画布，最终输出全零。

  • 排查：直接打印theta的数值，看是不是有远大于1或者远小于-1的离谱值。
  • 修复：给定位网络的输出加正则约束，比如在损失函数里加theta的L2正则；或者把最后一层激活换成tanh（把输出限制在[-1,1]），再乘个小系数（比如0.1）控制变换幅度，避免过度偏移。

• 采样网格超出有效范围
  PyTorch的grid_sample默认用padding_mode='zeros'，要是生成的采样网格所有坐标都落在输入图像的[-1,1]归一化范围外，输出自然全是0。

  • 排查：打印grid张量的数值范围，看是不是大部分坐标都超出了[-1,1]。
  • 修复：调整定位网络的输出范围，保证生成的变换不会把像素全移出去；也可以临时把padding_mode改成border或reflection验证，但核心还是得修正变换参数的生成逻辑。

• 梯度消失导致网络参数僵死
  要是定位网络的梯度消失，参数一直是初始化的随机值，大概率会生成无效的变换矩阵。

  • 排查：查看定位网络各层参数的grad值，要是梯度接近0，就是梯度消失了。
  • 修复：把ReLU换成LeakyReLU；简化定位网络的层数和神经元数量；调小学习率或者加梯度裁剪。

• 输入归一化不匹配
  输入图像的归一化范围和STN的采样逻辑不兼容，也可能导致采样异常。

  • 排查：确认输入张量的数值范围（比如是[0,1]还是[-1,1]），检查grid_sample的align_corners参数设置是否对应。
  • 修复：统一输入的归一化范围，比如归一化到[-1,1]，同时把align_corners设为True（多数STN示例的常规设置）。

快速验证方法 #

1. 注释掉STN的变换逻辑，直接返回输入张量，看输出是否正常，排除后续网络的干扰。
2. 把theta固定成单位矩阵（theta = torch.tensor([[1,0,0],[0,1,0]], device=device).unsqueeze(0).repeat(batch_size,1,1)），看输出是否和输入一致，验证grid_sample的采样逻辑没问题。
3. 缩小小批量大小，打印每一步的theta、grid和输出张量的数值，精准定位异常环节。

内容的提问来源于stack exchange，提问作者user20983853