Skip to content

refactor(metrics): extract classification helpers from Accuracy (#1222) #1429

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
michalharakal merged 2 commits into from
Oct 20, 2026
Merged

refactor(metrics): extract classification helpers from Accuracy (#1222) #1429

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
8b80ee5
refac: extracted helper function into ClassificationSupport
Oct 20, 2026
84851d3
refactor(metrics): decompose forEachPrediction in ClassificationSuppo…
Oct 20, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
208 changes: 4 additions & 204 deletions skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/metrics/Accuracy.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -36,23 +33,12 @@ public class Accuracy(
targets: Tensor<out DType, *>,
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<Int32, Int>, classDim)
}
FP32::class, FP16::class -> {
@Suppress("UNCHECKED_CAST")
computeWithSoftTargets(predictions, targets as Tensor<out DType, *>, 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 {
Expand All @@ -64,192 +50,6 @@ public class Accuracy(
correctCount = 0
totalCount = 0
}

private fun <T : DType, V> computeWithIndexTargets(
predictions: Tensor<T, V>,
targets: Tensor<Int32, Int>,
classDim: Int
): Pair<Long, Long> {
// 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 <T : DType, V> computeWithSoftTargets(
predictions: Tensor<T, V>,
targets: Tensor<out DType, *>,
classDim: Int
): Pair<Long, Long> {
// 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 <T : DType, V> argmax(
tensor: Tensor<T, V>,
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<out DType, *>,
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 <T : DType, V> validatePredictions(preds: Tensor<T, V>) {
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
}
}
}
}

/**
Expand Down
Loading
Loading