无法在C语言中获取TensorFlow Lite YOLO目标检测模型输出

我用Python训练了基于YOLO架构的TensorFlow Lite目标检测模型，现在要在C语言环境下执行推理，但调用解释器完成推理后，无法正确提取输出结果。

我的代码片段 #

模型加载与输入准备 #

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "tensorflow/lite/c/c_api.h"
#include "tensorflow/lite/c/c_api_experimental.h"
#include "tensorflow/lite/c/common.h"
#include "tensorflow/lite/c/builtin_op_data.h"

int inputImageHeight = 416;
int inputImageWidth = 416;
int channel = 3;
int imageSize = inputImageHeight * inputImageWidth * channel;
float imageDataBuffer[imageSize]; // 已将图像缩放到416x416并归一化到[0,1]，存储到该缓冲区

// 加载模型
TfLiteModel* model = TfLiteModelCreateFromFile(modelFilePath);

// 创建解释器
TfLiteInterpreter* interpreter = TfLiteInterpreterCreate(model, NULL);

// 分配张量
if(TfLiteInterpreterAllocateTensors(interpreter) != kTfLiteOk)
{
  disposeTfLiteObjects(model, interpreter);
  return -1;
}

// 调整输入张量维度
int inputDims[4] = {1, inputImageHeight, inputImageWidth, channel};
if(TfLiteInterpreterResizeInputTensor(interpreter, 0, inputDims, 4)!= kTfLiteOk)
{
  disposeTfLiteObjects(model, interpreter);
  return -1;
}

// 重新分配张量
if(TfLiteInterpreterAllocateTensors(interpreter) != kTfLiteOk)
{
  disposeTfLiteObjects(model, interpreter);
  return -1;
}

推理执行 #

// 获取输入张量
TfLiteTensor* inputTensor = TfLiteInterpreterGetInputTensor(interpreter, 0);

// 将图像数据复制到输入张量
if(TfLiteTensorCopyFromBuffer(inputTensor, imageDataBuffer, imageSize * sizeof(float)) != kTfLiteOk)
{
  disposeTfLiteObjects(model, interpreter);
  return -1;
}

// 执行推理
if(TfLiteInterpreterInvoke(interpreter) != kTfLiteOk)
{
  disposeTfLiteObjects(model, interpreter);
  return -1;
}

当前遇到的问题 #

上述步骤均正常执行，但提取输出时遇到两个问题：

1. 模型在Python中有两个输出张量，维度分别为[1,13,13,18]和[1,26,26,18]，但在C语言中只获取到第一个输出张量，找不到第二个。
2. 第一个输出张量的data.f字段第一个值为-nan(0x400000)，无法得到有效的原始检测数据。

调试信息 #

通过gdb查看第一个输出张量的维度，确认与Python中的[1,13,13,18]一致：

(gdb) print outputTensor->dims[0]
$10 = {size = 4, data = 0x555555785e04}
(gdb) print outputTensor->dims[1]
$11 = {size = 1, data = 0x555555785e08}
(gdb) print outputTensor->dims[2]
$12 = {size = 13, data = 0x555555785e0c}
(gdb) print outputTensor->dims[3]
$13 = {size = 13, data = 0x555555785e10}
(gdb) print outputTensor->dims[4]
$14 = {size = 18, data = 0x555555785e14}

查看输出张量数据指针：

(gdb) print outputTensor->data
$3 = {i32 = 0x7ffff561c8c0, i64 = 0x7ffff561c8c0, f = 0x7ffff561c8c0, f16 = 0x7ffff561c8c0, 
  f64 = 0x7ffff561c8c0, raw = 0x7ffff561c8c0 "", raw_const = 0x7ffff561c8c0 "", 
  uint8 = 0x7ffff561c8c0 "", b = 0x7ffff561c8c0, i16 = 0x7ffff561c8c0, c64 = 0x7ffff561c8c0, 
  c128 = 0x7ffff561c8c0, int8 = 0x7ffff561c8c0 "", data = 0x7ffff561c8c0}

读取第一个浮点值：

(gdb) print outputTensor->data.f[0]
$5 = -nan(0x400000)

我的需求是获取模型的原始输出结果，以便后续进行后处理（如非极大值抑制），也可以接受直接得到检测结果的方案。

解决方案 #

1. 正确获取所有输出张量 #

首先需要确认模型的输出张量总数，再逐个获取：

// 获取输出张量总数
int output_count = TfLiteInterpreterGetOutputTensorCount(interpreter);
printf("Total output tensors: %d\n", output_count);

// 遍历所有输出张量
for (int i = 0; i < output_count; i++) {
    const TfLiteTensor* output_tensor = TfLiteInterpreterGetOutputTensor(interpreter, i);
    
    // 打印输出张量的维度信息
    printf("Output tensor %d dimensions: ", i);
    for (int j = 0; j < output_tensor->dims->size; j++) {
        printf("%d ", output_tensor->dims->data[j]);
    }
    printf("\n");
}

通过TfLiteInterpreterGetOutputTensorCount可以拿到所有输出张量的数量，之后用索引0和1就能分别获取两个YOLO输出张量。

2. 解决NaN问题 #

出现NaN通常是输入数据或模型转换环节有问题，按以下步骤排查：

• 检查输入数据有效性：确认imageDataBuffer是否完全初始化，图像数据是否正确缩放到416x416，归一化是否确实到了[0,1]范围，有没有出现超出范围的值或未初始化的内存。
• 验证输入张量类型：用TfLiteTensorType(inputTensor)检查输入张量的类型，如果模型是浮点型（kTfLiteFloat32），则输入必须是float数组；如果是量化模型（比如kTfLiteUInt8），则需要将输入数据转换为对应类型，否则会导致推理输出异常。
• 检查模型转换过程：Python导出TFLite模型时，确保没有错误配置，比如是否开启了量化、是否保留了所有输出节点。可以在Python中重新验证模型的输出是否正常，再重新导出。

3. 提取原始输出数据 #

确认输出张量是浮点型后，可以直接通过data.f指针读取原始数据，按YOLO的格式解析：

// 以第一个输出张量为例，提取所有数据
const TfLiteTensor* output_tensor = TfLiteInterpreterGetOutputTensor(interpreter, 0);
int batch = output_tensor->dims->data[0];
int grid_h = output_tensor->dims->data[1];
int grid_w = output_tensor->dims->data[2];
int box_attrs = output_tensor->dims->data[3];

// 遍历每个网格、每个锚框
for (int b = 0; b < batch; b++) {
    for (int y = 0; y < grid_h; y++) {
        for (int x = 0; x < grid_w; x++) {
            for (int a = 0; a < 3; a++) { // 每个网格3个锚框
                int idx = b * grid_h * grid_w * box_attrs + y * grid_w * box_attrs + x * box_attrs + a * 6;
                float x_center = output_tensor->data.f[idx];
                float y_center = output_tensor->data.f[idx+1];
                float width = output_tensor->data.f[idx+2];
                float height = output_tensor->data.f[idx+3];
                float conf = output_tensor->data.f[idx+4];
                float class_prob = output_tensor->data.f[idx+5];
                
                // 这里可以添加边界框转换、置信度过滤等后处理逻辑
                printf("Grid (%d,%d), Anchor %d: conf=%.4f, class_prob=%.4f\n", x, y, a, conf, class_prob);
            }
        }
    }
}

4. 后处理提示 #

拿到原始输出后，需要按YOLO的逻辑将网格坐标转换为图像绝对坐标，然后过滤低置信度的检测框，最后执行非极大值抑制（NMS）去除重复框，得到最终的检测结果。

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