scn.onnx scatterDense之前输出正常，之后输出全为零

Question

你好，看到你的nvidia CUDA-centerpoint上的解答，很佩服。 现在有一个问题想请教一下，基于KITTI数据集训练了一个centerpoint模型，３Ｄ backbone采用 VoxelBackBone8x (底层同样基于sparseconvolution)，参考nvidia项目，导出scn.onnx时，scatterDense之前输出正常，之后输出全为零，这是什么原因？我的环境运行nvidia官方项目，能得到正确的结果，用自己的模型就卡在scatterDense这个算子，　请问可能的原因是什么？

Answer 1

晚上好🌙🌙🌙
本答案参考通义千问

你提到的问题是：在将训练好的 CenterPoint 模型导出为 ONNX 格式后，scatterDense 之前的输出正常，但之后的输出全为零。这是一个典型的 ONNX 导出问题，尤其是在使用 SparseConvolution（如 VoxelBackBone8x）结构时非常常见。

---

🚨 问题分析

1. ONNX 不支持稀疏操作
   SparseConvolution 是基于 SparseTensor 的操作，而 ONNX 标准不支持这种稀疏数据结构。因此，在导出 ONNX 时，框架（如 PyTorch 或 MMDetection3D）会尝试将这些操作转换为稠密形式。

2. scatterDense 算子未被正确识别或实现
   在 ONNX 中，scatterDense 并不是一个标准算子。如果模型中使用了自定义的 scatterDense 实现（例如通过 PyTorch 自定义操作），那么在导出到 ONNX 时可能会失败或导致输出为零。

3. 模型导出方式不一致
   如果你使用的是官方项目（如 NVIDIA 提供的 CUDA-CenterPoint），其 ONNX 导出流程可能已经处理了 sparseconv 和 scatterDense 的兼容性问题；而你自己导出的模型可能没有进行同样的优化或适配。

---

✅ 可能的原因总结

| 原因 | 描述 | |------|------| | ONNX 不支持 SparseTensor | SparseConvolution 无法直接导出为 ONNX，需转为稠密 | | scatterDense 算子未被正确实现 | 自定义的 scatterDense 在 ONNX 中无法识别或执行 | | 导出配置不一致 | 官方项目已适配 ONNX，而你的模型可能未做相应处理 |

---

🔧 解决方案

1. 确保 ONNX 支持稀疏操作

• 使用 torchscript 转换模型（推荐）：

  model = torch.jit.script(model)
  torch.jit.save(model, "model_scripted.pt")
  

• 使用 onnxsim 优化 ONNX 模型：

  pip install onnxsim
  onnxsim input_model.onnx output_model.onnx
  

2. 检查 scatterDense 的实现是否支持 ONNX

• 如果你是自己实现的 scatterDense，请确保它符合以下条件：
  • 使用 torch.nn.Module 或 torch.autograd.Function
  • 使用 @torch.jit.script 注解
  • 不使用任何非 ONNX 兼容的操作（如 index_put_、scatter_ 等）

3. 使用 ONNX 自定义算子（Custom Op）

• 如果 scatterDense 是自定义操作，可以将其封装为 ONNX 自定义算子。
• 参考文档：
  • ONNX Custom Operator

4. 使用 MMDetection3D 提供的 ONNX 导出工具

• 如果你是基于 MMDetection3D 训练的模型，建议使用其提供的导出脚本：

  python tools/export_onnx.py configs/centerpoint/centerpoint_voxel01_second_secfpn_64x16_kitti-3d-car.py \
    --work-dir ./work_dirs/centerpoint \
    --export-onnx-path ./centerpoint.onnx
  

5. 验证 ONNX 模型是否正确

• 使用 onnxruntime 验证模型：

  import onnxruntime as ort
  sess = ort.InferenceSession("scn.onnx")
  inputs = {input_name: input_data}
  outputs = sess.run(None, inputs)
  print(outputs)
  

---

💡 代码示例（修正 scatterDense）

如果你自己实现了 scatterDense，请参考以下简化版本：

import torch
from torch import nn
from torch.nn import functional as F

class ScatterDense(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, features, indices, shape):
        # 确保 indices 是 LongTensor
        indices = indices.long()
        # 创建一个全零张量
        out = torch.zeros(shape, dtype=features.dtype, device=features.device)
        # 使用 scatter_add_ 将特征散列到指定位置
        out.scatter_add_(0, indices, features)
        return out

# 示例使用
features = torch.randn(10, 32)
indices = torch.randint(0, 100, (10,))
shape = (100, 32)
scatter_op = ScatterDense()
output = scatter_op(features, indices, shape)
print(output.shape)  # 应该是 (100, 32)

注意： 上述代码仅为示例，实际应用中应根据你的模型结构调整。

---

✅ 总结

| 步骤 | 内容 | |------|------| | 1 | 确保 ONNX 支持稀疏操作，使用 torchscript 和 onnxsim 优化 | | 2 | 检查 scatterDense 是否支持 ONNX，必要时封装为自定义算子 | | 3 | 使用 MMDetection3D 提供的 ONNX 导出工具 | | 4 | 使用 onnxruntime 验证导出模型的正确性 | | 5 | 修改 scatterDense 实现以适应 ONNX 兼容性 |

---

如有需要，我可以帮你进一步调试你的 scn.onnx 文件，或者提供更详细的 scatterDense 适配方案。欢迎继续提问！