TensorRT：Invalid axis in softmax layer

最新推荐文章于 2024-04-29 10:08:50 发布

微风❤水墨

最新推荐文章于 2024-04-29 10:08:50 发布

阅读量355

点赞数

CC 4.0 BY-SA版权

分类专栏：目标检测-实现

本文链接：https://blog.csdn.net/lilai619/article/details/113260748

目标检测-实现专栏收录该内容

23 篇文章 ¥19.90 ¥99.00

订阅专栏

超级会员免费看

在尝试序列化RepVGG-A0模型到TensorRT时遇到错误，问题源于softmax层的无效轴。通过检查prototxt文件发现Reshape层缺少对NCHW维度的明确指定。为解决此问题，只需更新Reshape层，完整指定所有维度参数。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

尝试了一下 RepVGG-A0 的TensorRT序列化，遇到了下面这个问题：

Caffe Parser: Invalid axis in softmax layer - TensorRT expects NCHW input. Negative axis is not supported in TensorRT, please use positive axis indexing
error parsing layer type Softmax

看了一下prototxt： softmax前面，只有 Reshape layer里面是没有显示的指定NCHW【-1表示是自己推理处理】。

所以，稍微调整一下Reshape layer，将缺省维度参数补全即可：

layer {
  name: "view1"
  type: "Reshape"
  bottom: "ave_pool_blob1"
  top: "view_blob1"
  reshape_param {
    shape {
      dim: 0
      dim: -1
      dim: 1
      dim: 1
    }
  }
}