skainet-io-onnx/sk.ainet.io.onnx/OnnxWeightLoader

OnnxWeightLoader

Utility for loading ONNX model initializers (weights) into SKaiNET Module parameters.

Features:

ONNX-specific tensor decoding (FLOAT, INT32, INT64)
Format-agnostic weight mapping via WeightMapper
Support for multiple tensor data types (all converted to Float)
Debug mode for troubleshooting weight mapping issues

Example usage:

val model = ModelProto.decodeFromByteArray(bytes)
val ctx = DirectCpuExecutionContext(phase = Phase.INFERENCE)
val loadResult = OnnxWeightLoader.loadInitializers(model, ctx)
val mapping = OnnxWeightLoader.applyWeights(module, loadResult.tensors)

Types

InitializerLoadResult

data class ~~InitializerLoadResult~~(val tensors: List<OnnxWeightLoader.InitTensor>, val skipped: List<String>)

Result of loading initializers from an ONNX model.

InitTensor

data class ~~InitTensor~~(val name: String, val isBias: Boolean, val shape: List<Int>, val tensor: Tensor<FP32, Float>)

Holds a decoded ONNX tensor ready for weight assignment.

Functions

applyWeights

fun <T : DType, V> applyWeights(module: Module<T, V>, tensors: List<WeightTensor<T, V>>): WeightMapper.MappingResult

Apply ONNX weights to a module using name-based matching with shape-based fallback.

applyWeightsFromInitTensors

fun <T : DType, V> ~~applyWeightsFromInitTensors~~(module: Module<T, V>, tensors: List<OnnxWeightLoader.InitTensor>): WeightMapper.MappingResult

Apply weights from legacy InitTensor list.

applyWeightsWithDebug

fun <T : DType, V> applyWeightsWithDebug(module: Module<T, V>, tensors: List<WeightTensor<T, V>>, debug: Boolean): WeightMapper.MappingResult

Apply weights with optional debug output.

decodeInitializer

fun decodeInitializer(tensor: TensorProto, ctx: ExecutionContext): Result<WeightTensor<FP32, Float>>

Decode a single ONNX TensorProto into a WeightTensor.

extractOnnxNameFromPath

fun extractOnnxNameFromPath(path: String): String?

Extract the most specific "model.X.Y.Z" pattern from a DSL module path. E.g., "Yolo8/model.22.cv2.0/model.22.cv2.0.0/Conv2d-0" -> "model.22.cv2.0.0"

loadInitializers

fun loadInitializers(model: ModelProto, ctx: ExecutionContext): WeightLoadResult<FP32, Float>

Load all initializers (weights) from an ONNX ModelProto.

loadInitializersLegacy

fun ~~loadInitializersLegacy~~(model: ModelProto, ctx: ExecutionContext): OnnxWeightLoader.InitializerLoadResult

Load initializers returning the legacy InitializerLoadResult type.

matchesOnnxLayer

fun matchesOnnxLayer(onnxName: String, dslLayerName: String?): Boolean

Check if an ONNX tensor name matches a DSL layer name.

normalizeToOnnxFormat

fun normalizeToOnnxFormat(name: String): String

Normalize a DSL layer name to ONNX naming format. DSL uses underscores in some places while ONNX uses dots.

validateAllParametersMapped

fun validateAllParametersMapped(mapping: WeightMapper.MappingResult, skipped: List<String> = emptyList())

Validate that all module parameters were mapped from ONNX initializers. Throws IllegalArgumentException if mapping is incomplete.