ncnn 和 rknn 自定义算子对比实现

你提供了两种实现自定义 Sigmoid 算子的代码。第一种是使用NCNN 并行化的实现，第二种是一个自定义算子函数用于 RKNN（Rockchip Neural Network）的实现。下面我将详细解释这两种实现，并提供一些优化建议。

实现 1：使用NCNN 并行化实现的 Sigmoid 算子

int Sigmoid::forward_inplace(Mat& bottom_top_blob, const Option& opt) const
{
    int w = bottom_top_blob.w; // 获取张量的宽度
    int h = bottom_top_blob.h; // 获取张量的高度
    int d = bottom_top_blob.d; // 获取张量的深度
    int channels = bottom_top_blob.c; // 获取张量的通道数
    int size = w * h * d; // 计算每个通道的元素个数

    // 使用 OpenMP 并行化处理每个通道
    #pragma omp parallel for num_threads(opt.num_threads)
    for (int q = 0; q < channels; q++)
    {
        float* ptr = bottom_top_blob.channel(q); // 获取当前通道的数据指针

        for (int i = 0; i < size; i++)
        {
            float v = ptr[i]; // 获取当前元素值
            v = std::min(v, 88.3762626647949f); // 限制最大值以避免数值溢出
            v = std::max(v, -88.3762626647949f); // 限制最小值以避免数值溢出
            ptr[i] = 1.f / (1.f + expf(-v)); // 计算 Sigmoid 并存储结果
        }
    }

    return 0; // 返回 0 表示成功
}

解释

1. 获取张量维度：首先获取张量的宽度、高度、深度和通道数，并计算每个通道的元素个数。
2. OpenMP 并行化：使用 #pragma omp parallel for 指令并行化处理每个通道。
3. Sigmoid 计算：在每个通道中，对每个元素进行 Sigmoid 计算，并限制输入值的范围以避免数值溢出。

优化建议

• 减少重复计算：可以将常量 88.3762626647949f 提取出来，避免在循环中重复定义。
• 使用 SIMD 指令：如果编译器支持，可以使用 SIMD 指令（如 AVX）来加速计算。

实现 2：RKNN 自定义算子函数

int compute_custom_sigmoid_float32(rknn_custom_op_context* op_ctx, rknn_custom_op_tensor* inputs, uint32_t n_inputs,
                                    rknn_custom_op_tensor* outputs, uint32_t n_outputs)
{
    // 获取输入和输出张量的虚拟地址
    unsigned char* in_ptr = (unsigned char*)inputs[0].mem.virt_addr + inputs[0].mem.offset;
    unsigned char* out_ptr = (unsigned char*)outputs[0].mem.virt_addr + outputs[0].mem.offset;
    const float* in_data = (const float*)in_ptr; // 输入数据指针
    float* out_data = (float*)out_ptr; // 输出数据指针

    // 自定义 Sigmoid 算子的实现
    {
        int inside = 1; // 初始化元素个数

        // 计算输入张量的总元素个数
        for (int i = 0; i < inputs[0].attr.n_dims; i++) {
            inside *= inputs[0].attr.dims[i];
        }

        // 对每个元素进行 Sigmoid 计算
        for (int y = 0; y < inside; y++) {
            const float* src_y = in_data + y; // 获取当前元素的输入指针
            float* dst_y = out_data + y; // 获取当前元素的输出指针
            dst_y[0] = 1 / (1 + exp(-src_y[0])); // 计算 Sigmoid 并存储结果
        }
    }
    return 0; // 返回 0 表示成功
}

解释

1. 获取输入输出指针：从 rknn_custom_op_tensor 结构中获取输入和输出张量的虚拟地址。
2. 计算张量元素个数：通过输入张量的维度计算元素个数。
3. Sigmoid 计算：对每个元素进行 Sigmoid 计算，并将结果存储到输出张量中。

优化建议

• 减少重复计算：同样可以将常量提取出来，避免在循环中重复定义。
• 使用 SIMD 指令：可以使用 SIMD 指令来加速计算。

结论

这两种实现方式各有优劣。第一种使用 OpenMP 并行化，适合在多核 CPU 上运行；第二种是 RKNN 自定义算子实现，适合在 Rockchip 的 NPU 上运行。根据你的实际应用场景选择合适的实现方式，并结合上述优化建议进行性能提升。