YOLOv5转ptl模型推理时类别索引越界问题求助

问题概述 #

将自定义17类的YOLOv5模型通过Python转为PTL格式后，在Kotlin应用中加载推理，得到的预测类别索引为344938，远超出17类范围，触发数组越界异常：

java.lang.RuntimeException: Unable to start activity ComponentInfo{com.example.app/com.example.app.ImageInferenceActivity}: java.lang.ArrayIndexOutOfBoundsException: length=17; index=344938

转换与推理代码如下：

Python转换代码 #

import torch
from torch.utils.mobile_optimizer import optimize_for_mobile

torchscript_model = "best.torchscript"
export_model_name = "converted.ptl"

model = torch.jit.load(torchscript_model)
optimized_model = optimize_for_mobile(model)
optimized_model._save_for_lite_interpreter(export_model_name)

print(f"mobile optimized model exported to {export_model_name}")

Kotlin推理代码 #

val module = Module.load(assetFilePath(this, "converted.ptl"))
val bitmap = BitmapFactory.decodeStream(assets.open("1.jpg"))
val inputTensor = TensorImageUtils.bitmapToFloat32Tensor(
    bitmap,
    TensorImageUtils.TORCHVISION_NORM_MEAN_RGB,
    TensorImageUtils.TORCHVISION_NORM_STD_RGB
)
val output = module.forward(IValue.from(inputTensor))
val outputTensor = if (output.isTuple) {
    output.toTuple()[0].toTensor()
} else {
    output.toTensor()
}
val scores = outputTensor.dataAsFloatArray
var maxScoreIndex = 0
var maxScore = scores[0]

for (i in 1 until scores.size) {
    if (scores[i] > maxScore) {
        maxScore = scores[i]
        maxScoreIndex = i
    }
}
Log.d("prediction", "Predicted class: $maxScoreIndex")
val className: String = ImageNetClasses.IMAGENET_CLASSES[maxScoreIdx]

class ImageNetClasses {
    companion object {
        val IMAGENET_CLASSES = listOf(
            "100kph Maximum Speed Limit",
            "20T",
            "4.27m-Height-Limit",
            "40kph-Maximum-Speed-Limit",
            "50kph-Maximum-Speed-Limit",
            "60kph-Minimum-Speed-Limit",
            "Approaching-to-Intersection-Merging-Traffic",
            "Divided-Road-Ahead",
            "Give-Way",
            "Half-Y-Junction",
            "Merging-Traffic",
            "No-Overtaking",
            "No-Parking",
            "No-U-Turn",
            "Pass-Either-Side",
            "Road-Narrows-Left",
            "Turn-Right-Ahead",
        )
    }
}

核心原因 #

1. YOLOv5输出格式误解：YOLOv5是目标检测模型，输出并非分类模型的单类别概率数组，而是形状为[batch_size, num_anchors, 5 + num_classes]的张量。其中：

   • 前5个值：检测框的x, y, w, h坐标 + 框的置信度
   • 后17个值：对应自定义17类的类别概率
     当前将整个输出数组视为分类得分取最大索引，完全错误解析了输出结构，导致索引值异常。

2. 输入预处理不匹配：使用了ImageNet的归一化均值/方差（TORCHVISION_NORM_MEAN_RGB），但YOLOv5训练时的预处理是将图像像素值除以255，归一化均值为[0,0,0]、方差为[1,1,1]，预处理不匹配会导致输出结果失真。

3. 输入尺寸不匹配：YOLOv5默认输入尺寸为640x640，若输入Bitmap未缩放到该尺寸，会导致输出张量形状异常，进一步引发索引计算错误。

解决方案 #

1. 确保YOLOv5模型导出正确 #

先通过YOLOv5官方脚本导出torchscript模型，避免自定义导出可能的格式问题：

# 从YOLOv5根目录执行
python export.py --weights best.pt --include torchscript --img 640

得到best.torchscript后，再用原Python代码转为PTL格式。

2. 修改Kotlin输入预处理 #

调整归一化参数，同时将Bitmap缩放到640x640：

// 缩放Bitmap到YOLOv5输入尺寸640x640
val scaledBitmap = Bitmap.createScaledBitmap(bitmap, 640, 640, true)
// 使用YOLOv5匹配的归一化参数
val inputTensor = TensorImageUtils.bitmapToFloat32Tensor(
    scaledBitmap,
    floatArrayOf(0.0f, 0.0f, 0.0f),
    floatArrayOf(1.0f, 1.0f, 1.0f)
)

3. 正确解析YOLOv5输出张量 #

遍历每个检测框，过滤低置信度结果后提取类别索引：

val output = module.forward(IValue.from(inputTensor))
val outputTensor = if (output.isTuple) {
    output.toTuple()[0].toTensor()
} else {
    output.toTensor()
}
val scores = outputTensor.dataAsFloatArray
val shape = outputTensor.shape()
val numBoxes = shape[1] // 检测框数量
val numClasses = 17 // 自定义类别数
val confidenceThreshold = 0.5f // 置信度阈值，可调整

for (boxIdx in 0 until numBoxes) {
    // 获取当前框的置信度（第5个元素，索引从0开始）
    val confidence = scores[boxIdx * (5 + numClasses) + 4]
    // 过滤低置信度框
    if (confidence < confidenceThreshold) continue
    
    // 寻找当前框的最高概率类别
    var maxClassScore = 0f
    var maxClassIdx = 0
    for (classIdx in 0 until numClasses) {
        val classScore = scores[boxIdx * (5 + numClasses) + 5 + classIdx]
        if (classScore > maxClassScore) {
            maxClassScore = classScore
            maxClassIdx = classIdx
        }
    }
    
    // 输出有效预测结果
    val className = ImageNetClasses.IMAGENET_CLASSES[maxClassIdx]
    Log.d("prediction", "检测框$boxIdx: 类别=$className (索引=$maxClassIdx), 置信度=$confidence")
}

验证要点 #

• 确认输入Bitmap缩放至640x640，与模型训练时的输入尺寸一致
• 检查归一化参数与训练代码中的预处理逻辑完全匹配
• 输出解析时严格按照[x,y,w,h,conf, class1, class2, ..., class17]的结构提取数据

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