From 8b80ee563b612d6596868aef4f78e70a7c8a435e Mon Sep 17 00:00:00 2001 From: Edwin Date: Tue, 20 Oct 2026 14:29:06 +0200 Subject: [PATCH 1/2] refac: extracted helper function into ClassificationSupport - accuracy calls central "forEachPrediction" method for accumulation - helper functions are no longer private --- .../sk/ainet/lang/nn/metrics/Accuracy.kt | 205 +----------------- .../lang/nn/metrics/ClassificationSupport.kt | 175 +++++++++++++++ 2 files changed, 179 insertions(+), 201 deletions(-) create mode 100644 skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/metrics/ClassificationSupport.kt diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/metrics/Accuracy.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/metrics/Accuracy.kt index 7d0bb1602..ba2e4f0cb 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/metrics/Accuracy.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/metrics/Accuracy.kt @@ -36,23 +36,12 @@ public class Accuracy( targets: Tensor, ctx: ExecutionContext ) { - validatePredictions(predictions) - val classDim = normalizeDim(dim, predictions.rank) - - val (batchCorrect, batchTotal) = when (targets.dtype) { - Int32::class -> { - @Suppress("UNCHECKED_CAST") - computeWithIndexTargets(predictions, targets as Tensor, classDim) - } - FP32::class, FP16::class -> { - @Suppress("UNCHECKED_CAST") - computeWithSoftTargets(predictions, targets as Tensor, classDim) + ClassificationSupport.forEachPrediction(predictions, targets, dim, threshold) { predicted, target -> + if (predicted == target) { + correctCount++ } - else -> error("Unsupported target dtype for Accuracy: ${targets.dtype}") + totalCount++ } - - correctCount += batchCorrect - totalCount += batchTotal } override fun compute(): Double { @@ -64,192 +53,6 @@ public class Accuracy( correctCount = 0 totalCount = 0 } - - private fun computeWithIndexTargets( - predictions: Tensor, - targets: Tensor, - classDim: Int - ): Pair { - // Validate shapes: targets should have one less dimension than predictions - // (the class dimension is removed) - require(targets.rank == predictions.rank - 1) { - "Accuracy expected target rank ${predictions.rank - 1} for class indices, got ${targets.rank}" - } - validateIndexTargetShapes(predictions, targets, classDim) - - val classCount = predictions.shape[classDim] - var correct = 0L - var total = 0L - - // Iterate over all samples - iterateOverSamples(targets.shape.dimensions) { sampleIdx -> - val targetClass = targets.data.get(*sampleIdx) - val predictedClass = if (threshold != null && classCount <= 2) { - // Binary classification with threshold - val logitIdx = insertClassIndex(sampleIdx, if (classCount == 2) 1 else 0, classDim, predictions.rank) - val logit = predictions.data.get(*logitIdx) as Float - if (logit > threshold) 1 else 0 - } else { - // Multi-class: argmax along class dimension - argmax(predictions, sampleIdx, classDim) - } - - if (predictedClass == targetClass) { - correct++ - } - total++ - } - - return correct to total - } - - private fun computeWithSoftTargets( - predictions: Tensor, - targets: Tensor, - classDim: Int - ): Pair { - // Soft targets should have the same shape as predictions - require(predictions.shape == targets.shape) { - "Accuracy with soft targets requires matching shapes, got ${predictions.shape.dimensions.contentToString()} vs ${targets.shape.dimensions.contentToString()}" - } - - val classCount = predictions.shape[classDim] - var correct = 0L - var total = 0L - - // Get shape without the class dimension for iteration - val sampleShape = predictions.shape.dimensions.filterIndexed { i, _ -> i != classDim }.toIntArray() - - iterateOverSamples(sampleShape) { sampleIdx -> - val predictedClass = if (threshold != null && classCount <= 2) { - val logitIdx = insertClassIndex(sampleIdx, if (classCount == 2) 1 else 0, classDim, predictions.rank) - val logit = predictions.data.get(*logitIdx) as Float - if (logit > threshold) 1 else 0 - } else { - argmax(predictions, sampleIdx, classDim) - } - - val targetClass = argmaxTarget(targets, sampleIdx, classDim) - - if (predictedClass == targetClass) { - correct++ - } - total++ - } - - return correct to total - } - - private fun argmax( - tensor: Tensor, - sampleIdx: IntArray, - classDim: Int - ): Int { - val classCount = tensor.shape[classDim] - var maxIdx = 0 - var maxVal = Float.NEGATIVE_INFINITY - - for (c in 0 until classCount) { - val fullIdx = insertClassIndex(sampleIdx, c, classDim, tensor.rank) - val value = tensor.data.get(*fullIdx) as Float - if (value > maxVal) { - maxVal = value - maxIdx = c - } - } - return maxIdx - } - - private fun argmaxTarget( - tensor: Tensor, - sampleIdx: IntArray, - classDim: Int - ): Int { - val classCount = tensor.shape[classDim] - var maxIdx = 0 - var maxVal = Float.NEGATIVE_INFINITY - - for (c in 0 until classCount) { - val fullIdx = insertClassIndex(sampleIdx, c, classDim, tensor.rank) - val value = when (tensor.dtype) { - FP32::class -> tensor.data.get(*fullIdx) as Float - FP16::class -> (tensor.data.get(*fullIdx) as Number).toFloat() - else -> error("Unsupported dtype") - } - if (value > maxVal) { - maxVal = value - maxIdx = c - } - } - return maxIdx - } - - private fun insertClassIndex( - baseIdx: IntArray, - cls: Int, - classDim: Int, - outRank: Int - ): IntArray { - val result = IntArray(outRank) - var bi = 0 - for (i in 0 until outRank) { - if (i == classDim) { - result[i] = cls - } else { - result[i] = baseIdx[bi++] - } - } - return result - } - - private fun validateIndexTargetShapes( - preds: Tensor<*, *>, - targets: Tensor<*, *>, - classDim: Int - ) { - val predDims = preds.shape.dimensions - val targetDims = targets.shape.dimensions - var ti = 0 - for (pi in predDims.indices) { - if (pi == classDim) continue - require(predDims[pi] == targetDims[ti]) { - "Accuracy target shape mismatch at dim $pi: expected ${predDims[pi]}, got ${targetDims[ti]}" - } - ti++ - } - } - - private fun validatePredictions(preds: Tensor) { - require(preds.dtype == FP32::class || preds.dtype == FP16::class) { - "Accuracy requires floating point predictions, got ${preds.dtype}" - } - } - - private fun normalizeDim(dim: Int, rank: Int): Int { - val nd = if (dim < 0) dim + rank else dim - require(nd in 0 until rank) { "Dimension $dim out of range for rank $rank" } - return nd - } - - private inline fun iterateOverSamples(shape: IntArray, action: (IntArray) -> Unit) { - if (shape.isEmpty()) { - action(IntArray(0)) - return - } - - val idx = IntArray(shape.size) - val total = shape.fold(1) { acc, d -> acc * d } - - repeat(total) { - action(idx) - // Increment index - for (d in shape.lastIndex downTo 0) { - idx[d]++ - if (idx[d] < shape[d]) break - idx[d] = 0 - } - } - } } /** diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/metrics/ClassificationSupport.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/metrics/ClassificationSupport.kt new file mode 100644 index 000000000..85fc4f64e --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/metrics/ClassificationSupport.kt @@ -0,0 +1,175 @@ +package sk.ainet.lang.nn.metrics + +internal object ClassificationSupport { + + internal inline fun forEachPrediction( + predictions: Tensor, + targets: Tensor, + dim: Int = -1, + threshold: Float? = null, + action: (predicted: Int, target: Int) -> Unit + ){ + validatePredictions(predictions) + val classDim = normalizeDim(dim, predictions.rank) + + when (targets.dtype) { + Int32::class -> { + require(targets.rank == predictions.rank - 1) { + "Expected target rank ${predictions.rank - 1} for class indices, got ${targets.rank}" + } + validateIndexTargetShapes(predictions, targets, classDim) + + val classCount = predictions.shape[classDim] + @Suppress("UNCHECKED_CAST") + val intTargets = targets as Tensor + + iterateOverSamples(intTargets.shape.dimensions) { sampleIdx -> + val targetClass = intTargets.data.get(*sampleIdx) + val predictedClass = if (threshold != null && classCount <= 2) { + val logitIdx = insertClassIndex(sampleIdx, if (classCount == 2) 1 else 0, classDim, predictions.rank) + val logit = predictions.data.get(*logitIdx) as Float + if (logit > threshold) 1 else 0 + } else { + argmax(predictions, sampleIdx, classDim) + } + + action(predictedClass, targetClass) + } + } + FP32::class, FP16::class -> { + require(predictions.shape == targets.shape) { + "With soft targets requires matching shapes, got ${predictions.shape.dimensions.contentToString()} vs ${targets.shape.dimensions.contentToString()}" + } + + val classCount = predictions.shape[classDim] + val sampleShape = predictions.shape.dimensions.filterIndexed { i, _ -> i != classDim }.toIntArray() + + iterateOverSamples(sampleShape) { sampleIdx -> + val predictedClass = if (threshold != null && classCount <= 2) { + val logitIdx = insertClassIndex(sampleIdx, if (classCount == 2) 1 else 0, classDim, predictions.rank) + val logit = predictions.data.get(*logitIdx) as Float + if (logit > threshold) 1 else 0 + } else { + argmax(predictions, sampleIdx, classDim) + } + + val targetClass = argmaxTarget(targets, sampleIdx, classDim) + + action(predictedClass, targetClass) + } + } + else -> error("Unsupported target dtype for classification: ${targets.dtype}") + } + } + + internal fun argmax( + tensor: Tensor, + sampleIdx: IntArray, + classDim: Int + ): Int { + val classCount = tensor.shape[classDim] + var maxIdx = 0 + var maxVal = Float.NEGATIVE_INFINITY + + for (c in 0 until classCount) { + val fullIdx = insertClassIndex(sampleIdx, c, classDim, tensor.rank) + val value = tensor.data.get(*fullIdx) as Float + if (value > maxVal) { + maxVal = value + maxIdx = c + } + } + return maxIdx + } + + internal fun argmaxTarget( + tensor: Tensor, + sampleIdx: IntArray, + classDim: Int + ): Int { + val classCount = tensor.shape[classDim] + var maxIdx = 0 + var maxVal = Float.NEGATIVE_INFINITY + + for (c in 0 until classCount) { + val fullIdx = insertClassIndex(sampleIdx, c, classDim, tensor.rank) + val value = when (tensor.dtype) { + FP32::class -> tensor.data.get(*fullIdx) as Float + FP16::class -> (tensor.data.get(*fullIdx) as Number).toFloat() + else -> error("Unsupported dtype") + } + if (value > maxVal) { + maxVal = value + maxIdx = c + } + } + return maxIdx + } + + internal fun insertClassIndex( + baseIdx: IntArray, + cls: Int, + classDim: Int, + outRank: Int + ): IntArray { + val result = IntArray(outRank) + var bi = 0 + for (i in 0 until outRank) { + if (i == classDim) { + result[i] = cls + } else { + result[i] = baseIdx[bi++] + } + } + return result + } + + internal fun validateIndexTargetShapes( + preds: Tensor<*, *>, + targets: Tensor<*, *>, + classDim: Int + ) { + val predDims = preds.shape.dimensions + val targetDims = targets.shape.dimensions + var ti = 0 + for (pi in predDims.indices) { + if (pi == classDim) continue + require(predDims[pi] == targetDims[ti]) { + "Accuracy target shape mismatch at dim $pi: expected ${predDims[pi]}, got ${targetDims[ti]}" + } + ti++ + } + } + + internal fun validatePredictions(preds: Tensor) { + require(preds.dtype == FP32::class || preds.dtype == FP16::class) { + "Accuracy requires floating point predictions, got ${preds.dtype}" + } + } + + internal fun normalizeDim(dim: Int, rank: Int): Int { + val nd = if (dim < 0) dim + rank else dim + require(nd in 0 until rank) { "Dimension $dim out of range for rank $rank" } + return nd + } + + internal inline fun iterateOverSamples(shape: IntArray, action: (IntArray) -> Unit) { + if (shape.isEmpty()) { + action(IntArray(0)) + return + } + + val idx = IntArray(shape.size) + val total = shape.fold(1) { acc, d -> acc * d } + + repeat(total) { + action(idx) + // Increment index + for (d in shape.lastIndex downTo 0) { + idx[d]++ + if (idx[d] < shape[d]) break + idx[d] = 0 + } + } + } +} \ No newline at end of file From 84851d33495ac33e9528db36ae3cc1d443f4eeda Mon Sep 17 00:00:00 2001 From: Edwin Date: Tue, 20 Oct 2026 14:47:56 +0200 Subject: [PATCH 2/2] refactor(metrics): decompose forEachPrediction in ClassificationSupport (#1222) - Split forEachPrediction into forEachWithIndexTargets and forEachWithSoftTargets for clear separation between Int32 hard indices and FP32/FP16 soft targets. - Factor out duplicated sample traversal into iteratePredictions and class prediction resolution into predictClass. - Use inline and crossinline to maintain zero-allocation performance. - Remove obsolete DType imports from Accuracy.kt. --- .../sk/ainet/lang/nn/metrics/Accuracy.kt | 3 - .../lang/nn/metrics/ClassificationSupport.kt | 156 ++++++++++++------ 2 files changed, 105 insertions(+), 54 deletions(-) diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/metrics/Accuracy.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/metrics/Accuracy.kt index ba2e4f0cb..55bef4a97 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/metrics/Accuracy.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/metrics/Accuracy.kt @@ -3,9 +3,6 @@ package sk.ainet.lang.nn.metrics import sk.ainet.context.ExecutionContext import sk.ainet.lang.tensor.Tensor import sk.ainet.lang.types.DType -import sk.ainet.lang.types.FP16 -import sk.ainet.lang.types.FP32 -import sk.ainet.lang.types.Int32 /** * Classification accuracy metric. diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/metrics/ClassificationSupport.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/metrics/ClassificationSupport.kt index 85fc4f64e..119748a9a 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/metrics/ClassificationSupport.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/metrics/ClassificationSupport.kt @@ -1,67 +1,122 @@ package sk.ainet.lang.nn.metrics +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.types.DType +import sk.ainet.lang.types.FP16 +import sk.ainet.lang.types.FP32 +import sk.ainet.lang.types.Int32 + internal object ClassificationSupport { + /** + * Traverses predictions and targets sample by sample, invoking [action] with the + * predicted class and target class for each sample. + */ internal inline fun forEachPrediction( predictions: Tensor, targets: Tensor, dim: Int = -1, threshold: Float? = null, action: (predicted: Int, target: Int) -> Unit - ){ + ) { validatePredictions(predictions) val classDim = normalizeDim(dim, predictions.rank) when (targets.dtype) { - Int32::class -> { - require(targets.rank == predictions.rank - 1) { - "Expected target rank ${predictions.rank - 1} for class indices, got ${targets.rank}" - } - validateIndexTargetShapes(predictions, targets, classDim) - - val classCount = predictions.shape[classDim] - @Suppress("UNCHECKED_CAST") - val intTargets = targets as Tensor - - iterateOverSamples(intTargets.shape.dimensions) { sampleIdx -> - val targetClass = intTargets.data.get(*sampleIdx) - val predictedClass = if (threshold != null && classCount <= 2) { - val logitIdx = insertClassIndex(sampleIdx, if (classCount == 2) 1 else 0, classDim, predictions.rank) - val logit = predictions.data.get(*logitIdx) as Float - if (logit > threshold) 1 else 0 - } else { - argmax(predictions, sampleIdx, classDim) - } - - action(predictedClass, targetClass) - } - } - FP32::class, FP16::class -> { - require(predictions.shape == targets.shape) { - "With soft targets requires matching shapes, got ${predictions.shape.dimensions.contentToString()} vs ${targets.shape.dimensions.contentToString()}" - } - - val classCount = predictions.shape[classDim] - val sampleShape = predictions.shape.dimensions.filterIndexed { i, _ -> i != classDim }.toIntArray() - - iterateOverSamples(sampleShape) { sampleIdx -> - val predictedClass = if (threshold != null && classCount <= 2) { - val logitIdx = insertClassIndex(sampleIdx, if (classCount == 2) 1 else 0, classDim, predictions.rank) - val logit = predictions.data.get(*logitIdx) as Float - if (logit > threshold) 1 else 0 - } else { - argmax(predictions, sampleIdx, classDim) - } - - val targetClass = argmaxTarget(targets, sampleIdx, classDim) - - action(predictedClass, targetClass) - } - } + Int32::class -> forEachWithIndexTargets(predictions, targets, classDim, threshold, action) + FP32::class, FP16::class -> forEachWithSoftTargets(predictions, targets, classDim, threshold, action) else -> error("Unsupported target dtype for classification: ${targets.dtype}") } } + /** + * Handles traversal for hard class index targets (Int32). + */ + internal inline fun forEachWithIndexTargets( + predictions: Tensor, + targets: Tensor, + classDim: Int, + threshold: Float?, + action: (predicted: Int, target: Int) -> Unit + ) { + require(targets.rank == predictions.rank - 1) { + "Expected target rank ${predictions.rank - 1} for class indices, got ${targets.rank}" + } + validateIndexTargetShapes(predictions, targets, classDim) + + @Suppress("UNCHECKED_CAST") + val intTargets = targets as Tensor + iteratePredictions( + predictions = predictions, + sampleShape = intTargets.shape.dimensions, + classDim = classDim, + threshold = threshold, + getTargetClass = { sampleIdx -> intTargets.data.get(*sampleIdx) }, + action = action + ) + } + + /** + * Handles traversal for soft probability / one-hot targets (FP32 / FP16). + */ + internal inline fun forEachWithSoftTargets( + predictions: Tensor, + targets: Tensor, + classDim: Int, + threshold: Float?, + action: (predicted: Int, target: Int) -> Unit + ) { + require(predictions.shape == targets.shape) { + "With soft targets requires matching shapes, got ${predictions.shape.dimensions.contentToString()} vs ${targets.shape.dimensions.contentToString()}" + } + + val sampleShape = predictions.shape.dimensions.filterIndexed { i, _ -> i != classDim }.toIntArray() + iteratePredictions( + predictions = predictions, + sampleShape = sampleShape, + classDim = classDim, + threshold = threshold, + getTargetClass = { sampleIdx -> argmaxTarget(targets, sampleIdx, classDim) }, + action = action + ) + } + + /** + * Core sample traversal runner that coordinates per-sample prediction and target extraction. + */ + internal inline fun iteratePredictions( + predictions: Tensor, + sampleShape: IntArray, + classDim: Int, + threshold: Float?, + crossinline getTargetClass: (sampleIdx: IntArray) -> Int, + action: (predicted: Int, target: Int) -> Unit + ) { + val classCount = predictions.shape[classDim] + iterateOverSamples(sampleShape) { sampleIdx -> + val predictedClass = predictClass(predictions, sampleIdx, classDim, classCount, threshold) + val targetClass = getTargetClass(sampleIdx) + action(predictedClass, targetClass) + } + } + + /** + * Resolves the predicted class for a sample (thresholding for binary, argmax for multiclass). + */ + internal fun predictClass( + predictions: Tensor, + sampleIdx: IntArray, + classDim: Int, + classCount: Int, + threshold: Float? + ): Int = if (threshold != null && classCount <= 2) { + val logitIdx = insertClassIndex(sampleIdx, if (classCount == 2) 1 else 0, classDim, predictions.rank) + val logit = predictions.data.get(*logitIdx) as Float + if (logit > threshold) 1 else 0 + } else { + argmax(predictions, sampleIdx, classDim) + } + internal fun argmax( tensor: Tensor, sampleIdx: IntArray, @@ -96,7 +151,7 @@ internal object ClassificationSupport { val value = when (tensor.dtype) { FP32::class -> tensor.data.get(*fullIdx) as Float FP16::class -> (tensor.data.get(*fullIdx) as Number).toFloat() - else -> error("Unsupported dtype") + else -> error("Unsupported dtype: ${tensor.dtype}") } if (value > maxVal) { maxVal = value @@ -135,7 +190,7 @@ internal object ClassificationSupport { for (pi in predDims.indices) { if (pi == classDim) continue require(predDims[pi] == targetDims[ti]) { - "Accuracy target shape mismatch at dim $pi: expected ${predDims[pi]}, got ${targetDims[ti]}" + "Target shape mismatch at dim $pi: expected ${predDims[pi]}, got ${targetDims[ti]}" } ti++ } @@ -143,7 +198,7 @@ internal object ClassificationSupport { internal fun validatePredictions(preds: Tensor) { require(preds.dtype == FP32::class || preds.dtype == FP16::class) { - "Accuracy requires floating point predictions, got ${preds.dtype}" + "Requires floating point predictions, got ${preds.dtype}" } } @@ -164,7 +219,6 @@ internal object ClassificationSupport { repeat(total) { action(idx) - // Increment index for (d in shape.lastIndex downTo 0) { idx[d]++ if (idx[d] < shape[d]) break @@ -172,4 +226,4 @@ internal object ClassificationSupport { } } } -} \ No newline at end of file +}