ViViT视觉里程计回归测试阶段输出恒同问题求助

• 任务：基于KITTI里程计数据集的t与t+1时刻堆叠图像，预测智能体的位姿变化（位置+姿态）
• 模型：采用时空视觉Transformer（ViViT），输入张量形状为[B, C, F, H, W]（B=批量大小，C=通道数，F=堆叠帧数，H/W=图像高/宽）
• 异常表现：训练阶段模型输出多样，可准确拟合轨迹（蓝色为Ground Truth，橙色为预测轨迹）；但测试阶段所有输入（包括随机张量torch.rand(1, 3, 2, 370, 1241)）输出完全一致，无法预测位姿变化
• 训练配置：使用自定义损失函数与RMSProp优化器

• 检查模型状态与参数加载：
  • 确认测试时已调用model.eval()，关闭BatchNorm、Dropout等训练专属层的随机行为
  • 验证训练后保存的模型参数是否正确加载，可通过对比训练前后关键层的参数值（如注意力权重、回归头权重）确认
• 排查模型结构逻辑：
  • 重点检查时空注意力模块：确认注意力分数计算、权重归一化（如Softmax）逻辑是否正确，避免出现所有注意力权重一致的情况；检查模块参数初始化是否合理（如避免全零/全一初始化）
  • 检查最终回归头：确认参数未被误设置为requires_grad=False，激活函数使用符合回归任务需求（如避免误用会输出固定值的激活）
• 核对数据预处理流程：
  • 确保训练与测试阶段的图像预处理（如归一化、裁剪、维度转换）完全一致，避免输入分布差异导致模型输出异常
• 验证梯度传播有效性：
  • 在训练过程中监控各层的梯度值，确认梯度能正常传播到所有可训练层；若存在梯度消失/为零的层，需调整损失函数或模型结构
• 检查损失与优化器配置：
  • 验证自定义损失函数的梯度计算是否正确，可通过手动计算或使用torch.autograd.grad验证损失对参数的梯度是否非零
  • 检查RMSProp优化器的参数设置（如学习率、动量）是否合理，避免因学习率过低导致参数无法有效更新

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