J5 softmax 算子约束替换方法

在将深度学习模型部署到地平线J5开发板上时，我们遇到了一些与量化编译相关的挑战，尤其是与softmax操作的性能问题。在进行模型量化时，我们发现softmax中的一个关键算子——reducesum的维度超过了BPU（硬件加速单元）的限制，导致模型在推理时出现了分段现象，并且性能不如预期。本文将详细介绍如何解决这一问题，并提升模型推理性能。

问题描述

在进行量化编译时，softmax操作中涉及到一个reducesum算子，其维度在某些情况下可能会超过4096的限制，这在地平线J5开发板的BPU上是无法处理的。我们发现，当softmax中的reducesum算子维度过大时，该操作会被转移到CPU上执行，从而导致模型分段并显著降低推理速度。
Description

原因分析

softmax操作的典型过程包括：

ReduceMax：找到输入向量的最大值。
Exp：对每个元素进行指数运算。
ReduceSum：计算指数结果的总和。
Reciprocal：对总和取倒数。
Mul：将每个元素乘以倒数。

在地平线J5的BPU上，reducesum算子的输入维度如果超过4096，BPU无法有效处理该算子，会导致推理任务被转移到CPU上。这种情况显著影响了推理性能。

解决方法

为了解决这一问题，我们需要对softmax进行拆解，并对reducesum算子进行维度调整，使其符合BPU的硬件约束。

1. 拆解softmax操作

为了能对reducesum算子进行维度操作，我们需要先将将原本的softmax操作拆解成多个小算子。具体拆解步骤如下：

ReduceMax：计算最大值。
Sub：通过减法将每个输入减去最大值。
Exp：计算指数。
ReduceSum：对指数结果求和。
Reciprocal：取倒数。
Mul：将每个元素乘以倒数。

2. 调整reducesum的维度

首先，我们分析了reducesum算子的输入维度。输入的维度是1*8400*4*16，在进行reducesum时，这个维度会超过BPU的限制。为了解决这个问题，我们在reducesum算子之前插入了一个reshape算子，将输入维度从1*8400*4*16调整为1*2800*3*4*16。这样一来，reducesum的输出维度变为1*2800*3*4*1，从而满足了BPU对维度的约束。

3. 添加reshape操作

在reducesum之前和之后，我们分别插入了reshape算子。具体来说：

前置reshape：将输入的维度从1*8400*4*16调整为1*2800*3*4*16。
后置reshape：将reducesum的输出从1*2800*3*4*1转换回1*8400*4*1。

通过这种方式，我们可以有效避免reducesum算子维度过大而导致的问题，同时确保其运行在BPU上。

示例脚本

通过该脚本我们能够找到超出BPU约束的reducesum算子节点，并做出维度变换。

可视化效果

总结

在部署深度学习模型到地平线J5开发板时，面对softmax操作中的reducesum维度超限问题，我们通过插入reshape算子并拆解softmax操作，将所有算子都迁移到BPU上，从而避免了推理性能瓶颈。这一解决方案不仅提升了模型的推理效率，也让我们对硬件加速单元的约束有了更深入的理解。

在未来的项目中，我们将继续探索如何优化模型的量化过程，以便在各种硬件平台上都能获得更高效的推理性能。