[Python] TorchScript error: "The size of tensor a (3549) must match the size of tensor b...

I'm developing an Android application that uses a YOLOv8 model exported to TorchScript, which I load with PyTorch Mobile for object detection. The model was exported using Ultralytics with the following command:

# Cargar el modelo YOLOv8 preentrenado
model = YOLO("yolov8n.pt")

# Entrenar el modelo
model.train(data=yaml_path, epochs=50, imgsz=416, batch=8, workers=0, verbose=False, device=device)

# Guardar el modelo entrenado
model_path = "yolov8_fruits_model.pt"
model.export(format="torchscript", imgsz=640, simplify=True) # Guardar en formato TorchScript
print(f"Modelo guardado en: {model_path}")


However, when running the model on Android, I get the following error:

RuntimeError: The size of tensor a (3549) must match the size of tensor b (8400) at non-singleton dimension 2


Part of the stack trace points to:

torch.sub(anchor_points, lt) torch.add(anchor_points, rb)


This seems to be related to the decode_bboxes() function in head.py from Ultralytics. The error occurs when calling Module.forward(...) in Kotlin:

val output = module.forward(IValue.from(inputTensor)).toTensor()


What could be causing this error? Is it due to the input size or a model export issue? Is there a specific input size required or any additional step needed when exporting the model from Python?

Any guidance would be appreciated. Thanks!

pd. there is the following fun im using in my fragment:

private fun processImage(bitmap: Bitmap) {
if (yoloModule == null) {
showError("Modelo no cargado")
return
}

try {
val resizedBitmap = Bitmap.createScaledBitmap(bitmap, 640, 640, true) // Ajusta si tu modelo usa otro tamaño
val inputTensor = TensorImageUtils.bitmapToFloat32Tensor(
resizedBitmap,
floatArrayOf(0f, 0f, 0f), // Media RGB (ajusta si tu modelo usa otra normalización)
floatArrayOf(1f, 1f, 1f) // Std RGB
)

val output = yoloModule?.forward(IValue.from(inputTensor))
val detections = processOutput(output)

if (detections.isNotEmpty()) {
showDetectionsOnImage(bitmap, detections)
showMultipleDialogs(detections.sortedByDescending { it.confidence })
} else {
showError("No se detectaron productos")
}

} catch (e: Exception) {
showError("Error al procesar la imagen: ${e.message}")
e.printStackTrace()
}
}

private fun processOutput(output: IValue?): List<Detection> {
val detections = mutableListOf<Detection>()
if (output == null) return detections

try {
val outputTensor = output.toTensor()
val shape = outputTensor.shape()
Log.d("ModelOutput", "Forma del tensor: ${shape.toList()}") // Debug: Verifica la forma real

// Ajusta estos valores según la forma REAL de tu tensor
val numChannels = shape[1].toInt() // Debería ser 32 (4 coord + 28 clases)
val numAnchors = shape[2].toInt() // Debería ser 8400
val numClasses = numChannels - 4 // 28 clases

val outputData = outputTensor.dataAsFloatArray

for (i in 0 until numAnchors) {
val offset = i * numChannels

// Extraer coordenadas (xywh) - normalizadas [0-1]
val x = outputData[offset]
val y = outputData[offset + 1]
val w = outputData[offset + 2]
val h = outputData[offset + 3]

// Encontrar la clase con mayor probabilidad
var maxClassIndex = 0
var maxClassScore = 0f
for (j in 0 until numClasses) {
val score = outputData[offset + 4 + j]
if (score > maxClassScore) {
maxClassScore = score
maxClassIndex = j
}
}

// Umbral de confianza (ajusta según necesites)
if (maxClassScore < 0.5f) continue

// Convertir coordenadas normalizadas a píxeles (640x640)
val left = (x - w / 2) * 640
val top = (y - h / 2) * 640
val right = (x + w / 2) * 640
val bottom = (y + h / 2) * 640

val label = labels.getOrNull(maxClassIndex) ?: "Desconocido"
detections.add(
Detection(
label = label,
confidence = maxClassScore,
rect = RectF(left, top, right, bottom)
))
}
} catch (e: Exception) {
Log.e("ModelOutput", "Error al procesar la salida: ${e.message}")
}

return applyNms(detections)
}

private fun applyNms(detections: List<Detection>, iouThreshold: Float = 0.5f): List<Detection> {
val sortedDetections = detections.sortedByDescending { it.confidence }
val finalDetections = mutableListOf<Detection>()
val usedIndices = mutableSetOf<Int>()

for (i in sortedDetections.indices) {
if (usedIndices.contains(i)) continue

val currentDetection = sortedDetections
finalDetections.add(currentDetection)
usedIndices.add(i)

for (j in (i + 1) until sortedDetections.size) {
if (usedIndices.contains(j)) continue

val otherDetection = sortedDetections[j]
if (calculateIou(currentDetection.rect, otherDetection.rect) > iouThreshold) {
usedIndices.add(j)
}
}
}

return finalDetections
}

private fun calculateIou(rect1: RectF, rect2: RectF): Float {
val intersection = RectF()
intersection.setIntersect(rect1, rect2)

val intersectionArea = if (intersection.width() >= 0 && intersection.height() >= 0) {
intersection.width() * intersection.height()
} else {
0f
}

val unionArea = (rect1.width() * rect1.height()) + (rect2.width() * rect2.height()) - intersectionArea

return if (unionArea > 0) intersectionArea / unionArea else 0f
}

Continue reading...