1281 lines
48 KiB
Swift
1281 lines
48 KiB
Swift
@preconcurrency import COnnxRuntime
|
|
import Foundation
|
|
|
|
public enum ORTLoggingLevel: Int32 {
|
|
case verbose
|
|
case info
|
|
case warning
|
|
case error
|
|
case fatal
|
|
|
|
internal var native: OrtLoggingLevel {
|
|
OrtLoggingLevel(UInt32(rawValue))
|
|
}
|
|
}
|
|
|
|
public enum ORTValueType: Int32 {
|
|
case unknown
|
|
case tensor
|
|
|
|
internal init(native: ONNXType) throws {
|
|
switch native {
|
|
case ONNX_TYPE_UNKNOWN:
|
|
self = .unknown
|
|
case ONNX_TYPE_TENSOR:
|
|
self = .tensor
|
|
default:
|
|
throw ORTError.unsupported("Only tensor ONNX values are currently supported")
|
|
}
|
|
}
|
|
}
|
|
|
|
public enum ORTTensorElementDataType: Int32 {
|
|
case undefined
|
|
case float
|
|
case int8
|
|
case uInt8
|
|
case int32
|
|
case uInt32
|
|
case int64
|
|
case uInt64
|
|
case string
|
|
|
|
internal var native: ONNXTensorElementDataType {
|
|
switch self {
|
|
case .undefined:
|
|
return ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED
|
|
case .float:
|
|
return ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT
|
|
case .int8:
|
|
return ONNX_TENSOR_ELEMENT_DATA_TYPE_INT8
|
|
case .uInt8:
|
|
return ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8
|
|
case .int32:
|
|
return ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32
|
|
case .uInt32:
|
|
return ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT32
|
|
case .int64:
|
|
return ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64
|
|
case .uInt64:
|
|
return ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT64
|
|
case .string:
|
|
return ONNX_TENSOR_ELEMENT_DATA_TYPE_STRING
|
|
}
|
|
}
|
|
|
|
internal init(native: ONNXTensorElementDataType) throws {
|
|
switch native {
|
|
case ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED:
|
|
self = .undefined
|
|
case ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT:
|
|
self = .float
|
|
case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT8:
|
|
self = .int8
|
|
case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8:
|
|
self = .uInt8
|
|
case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32:
|
|
self = .int32
|
|
case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT32:
|
|
self = .uInt32
|
|
case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64:
|
|
self = .int64
|
|
case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT64:
|
|
self = .uInt64
|
|
case ONNX_TENSOR_ELEMENT_DATA_TYPE_STRING:
|
|
self = .string
|
|
default:
|
|
throw ORTError.unsupported("Unsupported tensor element type \(native.rawValue)")
|
|
}
|
|
}
|
|
}
|
|
|
|
public enum ORTGraphOptimizationLevel: Int32 {
|
|
case none
|
|
case basic
|
|
case extended
|
|
case layout
|
|
case all
|
|
|
|
internal var native: GraphOptimizationLevel {
|
|
switch self {
|
|
case .none:
|
|
return ORT_DISABLE_ALL
|
|
case .basic:
|
|
return ORT_ENABLE_BASIC
|
|
case .extended:
|
|
return ORT_ENABLE_EXTENDED
|
|
case .layout, .all:
|
|
return ORT_ENABLE_ALL
|
|
}
|
|
}
|
|
}
|
|
|
|
public typealias ORTCAPIRegisterCustomOpsFnPtr = UnsafeMutableRawPointer
|
|
|
|
private final class ORTValueStorage {
|
|
let handle: OpaquePointer
|
|
let retainedBuffer: NSMutableData?
|
|
|
|
init(handle: OpaquePointer, retainedBuffer: NSMutableData? = nil) {
|
|
self.handle = handle
|
|
self.retainedBuffer = retainedBuffer
|
|
}
|
|
|
|
deinit {
|
|
ORTNative.api.pointee.ReleaseValue?(handle)
|
|
}
|
|
}
|
|
|
|
public func ORTVersion() -> String? {
|
|
ORTNative.apiBase?.pointee.GetVersionString().map { String(cString: $0) }
|
|
}
|
|
|
|
public func ORTIsCoreMLExecutionProviderAvailable() -> Bool {
|
|
false
|
|
}
|
|
|
|
public func ORTSetSeed(_ seed: Int64) {
|
|
guard let trainingApi = ORTNative.trainingApi else {
|
|
return
|
|
}
|
|
|
|
_ = trainingApi.pointee.SetSeed?(seed)
|
|
}
|
|
|
|
public struct ORTValueTypeInfo {
|
|
public let type: ORTValueType
|
|
public let tensorTypeAndShapeInfo: ORTTensorTypeAndShapeInfo?
|
|
|
|
public init(type: ORTValueType, tensorTypeAndShapeInfo: ORTTensorTypeAndShapeInfo?) {
|
|
self.type = type
|
|
self.tensorTypeAndShapeInfo = tensorTypeAndShapeInfo
|
|
}
|
|
}
|
|
|
|
public struct ORTTensorTypeAndShapeInfo {
|
|
public let elementType: ORTTensorElementDataType
|
|
public let shape: [Int]
|
|
|
|
public init(elementType: ORTTensorElementDataType, shape: [Int]) {
|
|
self.elementType = elementType
|
|
self.shape = shape
|
|
}
|
|
}
|
|
|
|
public struct ORTCoreMLExecutionProviderOptions {
|
|
public var useCPUOnly: Bool
|
|
public var useCPUAndGPU: Bool
|
|
public var enableOnSubgraphs: Bool
|
|
public var onlyEnableForDevicesWithANE: Bool
|
|
public var onlyAllowStaticInputShapes: Bool
|
|
public var createMLProgram: Bool
|
|
|
|
public init(
|
|
useCPUOnly: Bool = false,
|
|
useCPUAndGPU: Bool = false,
|
|
enableOnSubgraphs: Bool = false,
|
|
onlyEnableForDevicesWithANE: Bool = false,
|
|
onlyAllowStaticInputShapes: Bool = false,
|
|
createMLProgram: Bool = false
|
|
) {
|
|
self.useCPUOnly = useCPUOnly
|
|
self.useCPUAndGPU = useCPUAndGPU
|
|
self.enableOnSubgraphs = enableOnSubgraphs
|
|
self.onlyEnableForDevicesWithANE = onlyEnableForDevicesWithANE
|
|
self.onlyAllowStaticInputShapes = onlyAllowStaticInputShapes
|
|
self.createMLProgram = createMLProgram
|
|
}
|
|
|
|
internal var providerOptions: [String: String] {
|
|
[
|
|
"use_cpu_only": ORTNative.boolString(useCPUOnly),
|
|
"use_cpu_and_gpu": ORTNative.boolString(useCPUAndGPU),
|
|
"enable_on_subgraphs": ORTNative.boolString(enableOnSubgraphs),
|
|
"only_enable_for_devices_with_ane": ORTNative.boolString(onlyEnableForDevicesWithANE),
|
|
"only_allow_static_input_shapes": ORTNative.boolString(onlyAllowStaticInputShapes),
|
|
"create_mlprogram": ORTNative.boolString(createMLProgram),
|
|
]
|
|
}
|
|
}
|
|
|
|
public struct ORTXnnpackExecutionProviderOptions {
|
|
public var intraOpNumThreads: Int32
|
|
|
|
public init(intraOpNumThreads: Int32 = 0) {
|
|
self.intraOpNumThreads = intraOpNumThreads
|
|
}
|
|
}
|
|
|
|
public final class ORTEnv {
|
|
public let loggingLevel: ORTLoggingLevel
|
|
internal let handle: OpaquePointer
|
|
|
|
public init(loggingLevel: ORTLoggingLevel) throws {
|
|
self.loggingLevel = loggingLevel
|
|
|
|
var env: OpaquePointer?
|
|
try ORTNative.check(
|
|
ORTNative.api.pointee.CreateEnv?(loggingLevel.native, "onnxruntime-swift-interface", &env),
|
|
action: "CreateEnv"
|
|
)
|
|
guard let env else {
|
|
throw ORTError.runtimeFailure(action: "CreateEnv", message: "Environment was not returned")
|
|
}
|
|
handle = env
|
|
}
|
|
|
|
deinit {
|
|
ORTNative.api.pointee.ReleaseEnv?(handle)
|
|
}
|
|
}
|
|
|
|
public final class ORTSession {
|
|
private let handle: OpaquePointer
|
|
|
|
public init(
|
|
env: ORTEnv,
|
|
modelPath: String,
|
|
sessionOptions: ORTSessionOptions? = nil
|
|
) throws {
|
|
var session: OpaquePointer?
|
|
let createdOptions = sessionOptions == nil ? try ORTSessionOptions() : nil
|
|
let optionsHandle = sessionOptions?.handle ?? createdOptions!.handle
|
|
|
|
try modelPath.withCString { modelPathPointer in
|
|
try ORTNative.check(
|
|
ORTNative.api.pointee.CreateSession?(env.handle, modelPathPointer, optionsHandle, &session),
|
|
action: "CreateSession"
|
|
)
|
|
}
|
|
|
|
guard let session else {
|
|
throw ORTError.runtimeFailure(action: "CreateSession", message: "Session was not returned")
|
|
}
|
|
handle = session
|
|
}
|
|
|
|
deinit {
|
|
ORTNative.api.pointee.ReleaseSession?(handle)
|
|
}
|
|
|
|
public func run(
|
|
inputs: [String: ORTValue],
|
|
outputs: [String: ORTValue],
|
|
runOptions: ORTRunOptions? = nil
|
|
) throws {
|
|
let sortedInputs = inputs.keys.sorted()
|
|
let sortedOutputs = outputs.keys.sorted()
|
|
let inputValues = sortedInputs.map { Optional(inputs[$0]!.handle) }
|
|
var outputValues = sortedOutputs.map { Optional(outputs[$0]!.handle) }
|
|
let outputCount = outputValues.count
|
|
|
|
try ORTNative.withCStringArray(sortedInputs) { inputNames in
|
|
try ORTNative.withCStringArray(sortedOutputs) { outputNames in
|
|
try inputNames.withUnsafeBufferPointer { inputNameBuffer in
|
|
try outputNames.withUnsafeBufferPointer { outputNameBuffer in
|
|
try inputValues.withUnsafeBufferPointer { inputValueBuffer in
|
|
try outputValues.withUnsafeMutableBufferPointer { outputValueBuffer in
|
|
try ORTNative.check(
|
|
ORTNative.api.pointee.Run?(
|
|
handle,
|
|
runOptions?.handle,
|
|
inputNameBuffer.baseAddress,
|
|
inputValueBuffer.baseAddress,
|
|
inputValues.count,
|
|
outputNameBuffer.baseAddress,
|
|
outputCount,
|
|
outputValueBuffer.baseAddress
|
|
),
|
|
action: "Run"
|
|
)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
public func run(
|
|
inputs: [String: ORTValue],
|
|
outputNames: Set<String>,
|
|
runOptions: ORTRunOptions? = nil
|
|
) throws -> [String: ORTValue] {
|
|
let sortedInputs = inputs.keys.sorted()
|
|
let sortedOutputs = outputNames.sorted()
|
|
|
|
let inputValues = sortedInputs.map { Optional(inputs[$0]!.handle) }
|
|
var outputValues = Array<OpaquePointer?>(repeating: nil, count: sortedOutputs.count)
|
|
let outputCount = sortedOutputs.count
|
|
|
|
try ORTNative.withCStringArray(sortedInputs) { inputNames in
|
|
try ORTNative.withCStringArray(sortedOutputs) { outputNamesArray in
|
|
try inputNames.withUnsafeBufferPointer { inputNameBuffer in
|
|
try outputNamesArray.withUnsafeBufferPointer { outputNameBuffer in
|
|
try inputValues.withUnsafeBufferPointer { inputValueBuffer in
|
|
try outputValues.withUnsafeMutableBufferPointer { outputValueBuffer in
|
|
try ORTNative.check(
|
|
ORTNative.api.pointee.Run?(
|
|
handle,
|
|
runOptions?.handle,
|
|
inputNameBuffer.baseAddress,
|
|
inputValueBuffer.baseAddress,
|
|
inputValues.count,
|
|
outputNameBuffer.baseAddress,
|
|
outputCount,
|
|
outputValueBuffer.baseAddress
|
|
),
|
|
action: "Run"
|
|
)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
var results: [String: ORTValue] = [:]
|
|
for (index, name) in sortedOutputs.enumerated() {
|
|
guard let valueHandle = outputValues[index] else {
|
|
throw ORTError.runtimeFailure(action: "Run", message: "Output \(name) was not returned")
|
|
}
|
|
results[name] = ORTValue(adopting: valueHandle)
|
|
}
|
|
return results
|
|
}
|
|
|
|
public func inputNames() throws -> [String] {
|
|
var count = 0
|
|
try ORTNative.check(ORTNative.api.pointee.SessionGetInputCount?(handle, &count), action: "SessionGetInputCount")
|
|
return try ORTNative.withAllocator { allocator in
|
|
var results: [String] = []
|
|
results.reserveCapacity(count)
|
|
for index in 0..<count {
|
|
var pointer: UnsafeMutablePointer<CChar>?
|
|
try ORTNative.check(
|
|
ORTNative.api.pointee.SessionGetInputName?(handle, index, allocator, &pointer),
|
|
action: "SessionGetInputName"
|
|
)
|
|
results.append(try ORTNative.copyAllocatedString(pointer, allocator: allocator))
|
|
}
|
|
return results
|
|
}
|
|
}
|
|
|
|
public func overridableInitializerNames() throws -> [String] {
|
|
var count = 0
|
|
try ORTNative.check(
|
|
ORTNative.api.pointee.SessionGetOverridableInitializerCount?(handle, &count),
|
|
action: "SessionGetOverridableInitializerCount"
|
|
)
|
|
return try ORTNative.withAllocator { allocator in
|
|
var results: [String] = []
|
|
results.reserveCapacity(count)
|
|
for index in 0..<count {
|
|
var pointer: UnsafeMutablePointer<CChar>?
|
|
try ORTNative.check(
|
|
ORTNative.api.pointee.SessionGetOverridableInitializerName?(handle, index, allocator, &pointer),
|
|
action: "SessionGetOverridableInitializerName"
|
|
)
|
|
results.append(try ORTNative.copyAllocatedString(pointer, allocator: allocator))
|
|
}
|
|
return results
|
|
}
|
|
}
|
|
|
|
public func outputNames() throws -> [String] {
|
|
var count = 0
|
|
try ORTNative.check(ORTNative.api.pointee.SessionGetOutputCount?(handle, &count), action: "SessionGetOutputCount")
|
|
return try ORTNative.withAllocator { allocator in
|
|
var results: [String] = []
|
|
results.reserveCapacity(count)
|
|
for index in 0..<count {
|
|
var pointer: UnsafeMutablePointer<CChar>?
|
|
try ORTNative.check(
|
|
ORTNative.api.pointee.SessionGetOutputName?(handle, index, allocator, &pointer),
|
|
action: "SessionGetOutputName"
|
|
)
|
|
results.append(try ORTNative.copyAllocatedString(pointer, allocator: allocator))
|
|
}
|
|
return results
|
|
}
|
|
}
|
|
}
|
|
|
|
public final class ORTSessionOptions {
|
|
public private(set) var executionProviders: [(name: String, options: [String: String])] = []
|
|
public private(set) var intraOpNumThreads: Int32 = 0
|
|
public private(set) var graphOptimizationLevel: ORTGraphOptimizationLevel = .all
|
|
public private(set) var optimizedModelFilePath: String?
|
|
public private(set) var logID: String?
|
|
public private(set) var logSeverityLevel: ORTLoggingLevel?
|
|
public private(set) var configEntries: [String: String] = [:]
|
|
public private(set) var customOpRegistrationFunctionName: String?
|
|
public private(set) var customOpRegistrationFunctionPointer: ORTCAPIRegisterCustomOpsFnPtr?
|
|
public private(set) var ortExtensionsCustomOpsEnabled = false
|
|
public private(set) var coreMLExecutionProviderOptions: ORTCoreMLExecutionProviderOptions?
|
|
public private(set) var coreMLExecutionProviderV2Options: [String: String]?
|
|
public private(set) var xnnpackExecutionProviderOptions: ORTXnnpackExecutionProviderOptions?
|
|
|
|
internal let handle: OpaquePointer
|
|
|
|
public init() throws {
|
|
var options: OpaquePointer?
|
|
try ORTNative.check(ORTNative.api.pointee.CreateSessionOptions?(&options), action: "CreateSessionOptions")
|
|
guard let options else {
|
|
throw ORTError.runtimeFailure(action: "CreateSessionOptions", message: "Session options were not returned")
|
|
}
|
|
handle = options
|
|
}
|
|
|
|
deinit {
|
|
ORTNative.api.pointee.ReleaseSessionOptions?(handle)
|
|
}
|
|
|
|
public func appendExecutionProvider(
|
|
_ providerName: String,
|
|
providerOptions: [String: String]
|
|
) throws {
|
|
let sorted = providerOptions.keys.sorted().map { ($0, providerOptions[$0]!) }
|
|
let keyStorage = sorted.map { Array($0.0.utf8CString) }
|
|
let valueStorage = sorted.map { Array($0.1.utf8CString) }
|
|
let keys = keyStorage.map { $0.withUnsafeBufferPointer { $0.baseAddress } }
|
|
let values = valueStorage.map { $0.withUnsafeBufferPointer { $0.baseAddress } }
|
|
|
|
try providerName.withCString { providerNamePointer in
|
|
try keys.withUnsafeBufferPointer { keyBuffer in
|
|
try values.withUnsafeBufferPointer { valueBuffer in
|
|
try ORTNative.check(
|
|
ORTNative.api.pointee.SessionOptionsAppendExecutionProvider?(
|
|
handle,
|
|
providerNamePointer,
|
|
keyBuffer.baseAddress,
|
|
valueBuffer.baseAddress,
|
|
sorted.count
|
|
),
|
|
action: "SessionOptionsAppendExecutionProvider"
|
|
)
|
|
}
|
|
}
|
|
}
|
|
|
|
executionProviders.append((providerName, providerOptions))
|
|
}
|
|
|
|
public func setIntraOpNumThreads(_ count: Int32) throws {
|
|
try ORTNative.check(
|
|
ORTNative.api.pointee.SetIntraOpNumThreads?(handle, Int32(count)),
|
|
action: "SetIntraOpNumThreads"
|
|
)
|
|
intraOpNumThreads = count
|
|
}
|
|
|
|
public func setGraphOptimizationLevel(_ level: ORTGraphOptimizationLevel) throws {
|
|
try ORTNative.check(
|
|
ORTNative.api.pointee.SetSessionGraphOptimizationLevel?(handle, level.native),
|
|
action: "SetSessionGraphOptimizationLevel"
|
|
)
|
|
graphOptimizationLevel = level
|
|
}
|
|
|
|
public func setOptimizedModelFilePath(_ path: String) throws {
|
|
try path.withCString { pointer in
|
|
try ORTNative.check(
|
|
ORTNative.api.pointee.SetOptimizedModelFilePath?(handle, pointer),
|
|
action: "SetOptimizedModelFilePath"
|
|
)
|
|
}
|
|
optimizedModelFilePath = path
|
|
}
|
|
|
|
public func setLogID(_ logID: String) throws {
|
|
try logID.withCString { pointer in
|
|
try ORTNative.check(
|
|
ORTNative.api.pointee.SetSessionLogId?(handle, pointer),
|
|
action: "SetSessionLogId"
|
|
)
|
|
}
|
|
self.logID = logID
|
|
}
|
|
|
|
public func setLogSeverityLevel(_ level: ORTLoggingLevel) throws {
|
|
try ORTNative.check(
|
|
ORTNative.api.pointee.SetSessionLogSeverityLevel?(handle, level.rawValue),
|
|
action: "SetSessionLogSeverityLevel"
|
|
)
|
|
logSeverityLevel = level
|
|
}
|
|
|
|
public func addConfigEntry(key: String, value: String) throws {
|
|
try key.withCString { keyPointer in
|
|
try value.withCString { valuePointer in
|
|
try ORTNative.check(
|
|
ORTNative.api.pointee.AddSessionConfigEntry?(handle, keyPointer, valuePointer),
|
|
action: "AddSessionConfigEntry"
|
|
)
|
|
}
|
|
}
|
|
configEntries[key] = value
|
|
}
|
|
|
|
public func registerCustomOps(usingFunction name: String) throws {
|
|
try name.withCString { pointer in
|
|
try ORTNative.check(
|
|
ORTNative.api.pointee.RegisterCustomOpsUsingFunction?(handle, pointer),
|
|
action: "RegisterCustomOpsUsingFunction"
|
|
)
|
|
}
|
|
customOpRegistrationFunctionName = name
|
|
}
|
|
|
|
public func registerCustomOps(usingFunctionPointer fn: ORTCAPIRegisterCustomOpsFnPtr) throws {
|
|
let nativeFunction = unsafeBitCast(fn, to: RegisterCustomOpsFn.self)
|
|
try ORTNative.check(
|
|
nativeFunction(handle, ORTNative.apiBase),
|
|
action: "RegisterCustomOpsUsingFunctionPointer"
|
|
)
|
|
customOpRegistrationFunctionPointer = fn
|
|
}
|
|
|
|
public func enableOrtExtensionsCustomOps() throws {
|
|
try ORTNative.check(
|
|
ORTNative.api.pointee.EnableOrtCustomOps?(handle),
|
|
action: "EnableOrtCustomOps"
|
|
)
|
|
ortExtensionsCustomOpsEnabled = true
|
|
}
|
|
|
|
public func appendCoreMLExecutionProvider(options: ORTCoreMLExecutionProviderOptions) throws {
|
|
try appendCoreMLExecutionProviderV2(providerOptions: options.providerOptions)
|
|
coreMLExecutionProviderOptions = options
|
|
}
|
|
|
|
public func appendCoreMLExecutionProviderV2(providerOptions: [String: String]) throws {
|
|
try appendExecutionProvider("CoreML", providerOptions: providerOptions)
|
|
coreMLExecutionProviderV2Options = providerOptions
|
|
}
|
|
|
|
public func appendXnnpackExecutionProvider(options: ORTXnnpackExecutionProviderOptions) throws {
|
|
try appendExecutionProvider(
|
|
"XNNPACK",
|
|
providerOptions: ["intra_op_num_threads": String(options.intraOpNumThreads)]
|
|
)
|
|
xnnpackExecutionProviderOptions = options
|
|
}
|
|
}
|
|
|
|
public final class ORTRunOptions {
|
|
public private(set) var logTag: String?
|
|
public private(set) var logSeverityLevel: ORTLoggingLevel?
|
|
public private(set) var configEntries: [String: String] = [:]
|
|
|
|
internal let handle: OpaquePointer
|
|
|
|
public init() throws {
|
|
var options: OpaquePointer?
|
|
try ORTNative.check(ORTNative.api.pointee.CreateRunOptions?(&options), action: "CreateRunOptions")
|
|
guard let options else {
|
|
throw ORTError.runtimeFailure(action: "CreateRunOptions", message: "Run options were not returned")
|
|
}
|
|
handle = options
|
|
}
|
|
|
|
deinit {
|
|
ORTNative.api.pointee.ReleaseRunOptions?(handle)
|
|
}
|
|
|
|
public func setLogTag(_ tag: String) throws {
|
|
try tag.withCString { pointer in
|
|
try ORTNative.check(
|
|
ORTNative.api.pointee.RunOptionsSetRunTag?(handle, pointer),
|
|
action: "RunOptionsSetRunTag"
|
|
)
|
|
}
|
|
logTag = tag
|
|
}
|
|
|
|
public func setLogSeverityLevel(_ level: ORTLoggingLevel) throws {
|
|
try ORTNative.check(
|
|
ORTNative.api.pointee.RunOptionsSetRunLogSeverityLevel?(handle, level.rawValue),
|
|
action: "RunOptionsSetRunLogSeverityLevel"
|
|
)
|
|
logSeverityLevel = level
|
|
}
|
|
|
|
public func addConfigEntry(key: String, value: String) throws {
|
|
try key.withCString { keyPointer in
|
|
try value.withCString { valuePointer in
|
|
try ORTNative.check(
|
|
ORTNative.api.pointee.AddRunConfigEntry?(handle, keyPointer, valuePointer),
|
|
action: "AddRunConfigEntry"
|
|
)
|
|
}
|
|
}
|
|
configEntries[key] = value
|
|
}
|
|
}
|
|
|
|
public final class ORTValue {
|
|
private let storage: ORTValueStorage
|
|
|
|
internal var handle: OpaquePointer {
|
|
storage.handle
|
|
}
|
|
|
|
public init(
|
|
tensorData: NSMutableData,
|
|
elementType: ORTTensorElementDataType,
|
|
shape: [Int]
|
|
) throws {
|
|
let elementCount = try ORTNative.tensorElementCount(shape: shape)
|
|
let expectedLength = try ORTNative.elementByteCount(for: elementType) * elementCount
|
|
guard tensorData.length >= expectedLength else {
|
|
throw ORTError.invalidArgument("Tensor buffer is too small for shape \(shape) and element type \(elementType)")
|
|
}
|
|
|
|
var memoryInfo: OpaquePointer?
|
|
try ORTNative.check(
|
|
ORTNative.api.pointee.CreateCpuMemoryInfo?(OrtArenaAllocator, OrtMemTypeDefault, &memoryInfo),
|
|
action: "CreateCpuMemoryInfo"
|
|
)
|
|
guard let memoryInfo else {
|
|
throw ORTError.runtimeFailure(action: "CreateCpuMemoryInfo", message: "Memory info was not returned")
|
|
}
|
|
defer {
|
|
ORTNative.api.pointee.ReleaseMemoryInfo?(memoryInfo)
|
|
}
|
|
|
|
var value: OpaquePointer?
|
|
var nativeShape = ORTNative.makeShapeBuffer(shape)
|
|
let nativeShapeCount = nativeShape.count
|
|
try nativeShape.withUnsafeMutableBufferPointer { shapeBuffer in
|
|
try ORTNative.check(
|
|
ORTNative.api.pointee.CreateTensorWithDataAsOrtValue?(
|
|
memoryInfo,
|
|
tensorData.mutableBytes,
|
|
tensorData.length,
|
|
shapeBuffer.baseAddress,
|
|
nativeShapeCount,
|
|
elementType.native,
|
|
&value
|
|
),
|
|
action: "CreateTensorWithDataAsOrtValue"
|
|
)
|
|
}
|
|
|
|
guard let value else {
|
|
throw ORTError.runtimeFailure(action: "CreateTensorWithDataAsOrtValue", message: "Value was not returned")
|
|
}
|
|
storage = ORTValueStorage(handle: value, retainedBuffer: tensorData)
|
|
}
|
|
|
|
public init(
|
|
tensorStringData: [String],
|
|
shape: [Int]
|
|
) throws {
|
|
let elementCount = try ORTNative.tensorElementCount(shape: shape)
|
|
guard elementCount == tensorStringData.count else {
|
|
throw ORTError.invalidArgument("String tensor element count does not match shape")
|
|
}
|
|
|
|
var value: OpaquePointer?
|
|
var nativeShape = ORTNative.makeShapeBuffer(shape)
|
|
let nativeShapeCount = nativeShape.count
|
|
try ORTNative.withAllocator { allocator in
|
|
try nativeShape.withUnsafeMutableBufferPointer { shapeBuffer in
|
|
try ORTNative.check(
|
|
ORTNative.api.pointee.CreateTensorAsOrtValue?(
|
|
allocator,
|
|
shapeBuffer.baseAddress,
|
|
nativeShapeCount,
|
|
ONNX_TENSOR_ELEMENT_DATA_TYPE_STRING,
|
|
&value
|
|
),
|
|
action: "CreateTensorAsOrtValue"
|
|
)
|
|
}
|
|
}
|
|
|
|
guard let value else {
|
|
throw ORTError.runtimeFailure(action: "CreateTensorAsOrtValue", message: "String tensor value was not returned")
|
|
}
|
|
|
|
let stringStorage = tensorStringData.map { Array($0.utf8CString) }
|
|
let stringPointers = stringStorage.map { $0.withUnsafeBufferPointer { $0.baseAddress } }
|
|
try stringPointers.withUnsafeBufferPointer { pointerBuffer in
|
|
try ORTNative.check(
|
|
ORTNative.api.pointee.FillStringTensor?(value, pointerBuffer.baseAddress, stringPointers.count),
|
|
action: "FillStringTensor"
|
|
)
|
|
}
|
|
|
|
storage = ORTValueStorage(handle: value)
|
|
}
|
|
|
|
internal init(adopting handle: OpaquePointer) {
|
|
storage = ORTValueStorage(handle: handle)
|
|
}
|
|
|
|
public func typeInfo() throws -> ORTValueTypeInfo {
|
|
try ORTNative.fetchValueTypeInfo(from: handle)
|
|
}
|
|
|
|
public func tensorTypeAndShapeInfo() throws -> ORTTensorTypeAndShapeInfo {
|
|
try ORTNative.fetchTensorInfo(from: handle)
|
|
}
|
|
|
|
public func tensorData() throws -> NSMutableData {
|
|
let tensorInfo = try tensorTypeAndShapeInfo()
|
|
guard tensorInfo.elementType != .string else {
|
|
throw ORTError.unsupported("tensorData() is not available for string tensors")
|
|
}
|
|
|
|
var dataPointer: UnsafeMutableRawPointer?
|
|
try ORTNative.check(
|
|
ORTNative.api.pointee.GetTensorMutableData?(handle, &dataPointer),
|
|
action: "GetTensorMutableData"
|
|
)
|
|
guard let dataPointer else {
|
|
throw ORTError.runtimeFailure(action: "GetTensorMutableData", message: "Tensor data pointer was not returned")
|
|
}
|
|
|
|
let elementCount = try ORTNative.tensorElementCount(shape: tensorInfo.shape)
|
|
let length = try ORTNative.elementByteCount(for: tensorInfo.elementType) * elementCount
|
|
return NSMutableData(bytesNoCopy: dataPointer, length: length, freeWhenDone: false)
|
|
}
|
|
|
|
public func tensorStringData() throws -> [String] {
|
|
let tensorInfo = try tensorTypeAndShapeInfo()
|
|
guard tensorInfo.elementType == .string else {
|
|
throw ORTError.unsupported("tensorStringData() is only available for string tensors")
|
|
}
|
|
|
|
let elementCount = try ORTNative.tensorElementCount(shape: tensorInfo.shape)
|
|
var totalLength = 0
|
|
try ORTNative.check(
|
|
ORTNative.api.pointee.GetStringTensorDataLength?(handle, &totalLength),
|
|
action: "GetStringTensorDataLength"
|
|
)
|
|
|
|
var content = Array(repeating: UInt8(0), count: totalLength)
|
|
var offsets = Array(repeating: Int(0), count: elementCount)
|
|
try ORTNative.check(
|
|
ORTNative.api.pointee.GetStringTensorContent?(
|
|
handle,
|
|
&content,
|
|
totalLength,
|
|
&offsets,
|
|
elementCount
|
|
),
|
|
action: "GetStringTensorContent"
|
|
)
|
|
|
|
return (0..<elementCount).map { index in
|
|
let start = offsets[index]
|
|
let end = index + 1 < offsets.count ? offsets[index + 1] : totalLength
|
|
return String(decoding: content[start..<end], as: UTF8.self)
|
|
}
|
|
}
|
|
}
|
|
|
|
public final class ORTCheckpoint {
|
|
internal let handle: OpaquePointer
|
|
|
|
public init(path: String) throws {
|
|
guard let trainingApi = ORTNative.trainingApi else {
|
|
throw ORTError.unsupported("This ONNX Runtime build does not expose the training API")
|
|
}
|
|
|
|
var checkpoint: OpaquePointer?
|
|
try path.withCString { pointer in
|
|
try ORTNative.check(
|
|
trainingApi.pointee.LoadCheckpoint?(pointer, &checkpoint),
|
|
action: "LoadCheckpoint"
|
|
)
|
|
}
|
|
|
|
guard let checkpoint else {
|
|
throw ORTError.runtimeFailure(action: "LoadCheckpoint", message: "Checkpoint state was not returned")
|
|
}
|
|
handle = checkpoint
|
|
}
|
|
|
|
deinit {
|
|
guard let trainingApi = ORTNative.trainingApi else {
|
|
return
|
|
}
|
|
trainingApi.pointee.ReleaseCheckpointState?(handle)
|
|
}
|
|
|
|
public func saveCheckpoint(
|
|
to path: String,
|
|
includeOptimizerState: Bool
|
|
) throws {
|
|
guard let trainingApi = ORTNative.trainingApi else {
|
|
throw ORTError.unsupported("This ONNX Runtime build does not expose the training API")
|
|
}
|
|
|
|
try path.withCString { pointer in
|
|
try ORTNative.check(
|
|
trainingApi.pointee.SaveCheckpoint?(handle, pointer, includeOptimizerState),
|
|
action: "SaveCheckpoint"
|
|
)
|
|
}
|
|
}
|
|
|
|
public func addIntProperty(name: String, value: Int64) throws {
|
|
guard let trainingApi = ORTNative.trainingApi else {
|
|
throw ORTError.unsupported("This ONNX Runtime build does not expose the training API")
|
|
}
|
|
|
|
var mutableValue = value
|
|
try name.withCString { pointer in
|
|
try ORTNative.check(
|
|
trainingApi.pointee.AddProperty?(handle, pointer, OrtIntProperty, &mutableValue),
|
|
action: "AddProperty"
|
|
)
|
|
}
|
|
}
|
|
|
|
public func addFloatProperty(name: String, value: Float) throws {
|
|
guard let trainingApi = ORTNative.trainingApi else {
|
|
throw ORTError.unsupported("This ONNX Runtime build does not expose the training API")
|
|
}
|
|
|
|
var mutableValue = value
|
|
try name.withCString { pointer in
|
|
try ORTNative.check(
|
|
trainingApi.pointee.AddProperty?(handle, pointer, OrtFloatProperty, &mutableValue),
|
|
action: "AddProperty"
|
|
)
|
|
}
|
|
}
|
|
|
|
public func addStringProperty(name: String, value: String) throws {
|
|
guard let trainingApi = ORTNative.trainingApi else {
|
|
throw ORTError.unsupported("This ONNX Runtime build does not expose the training API")
|
|
}
|
|
|
|
try name.withCString { namePointer in
|
|
try value.withCString { valuePointer in
|
|
let mutablePointer = UnsafeMutableRawPointer(mutating: valuePointer)
|
|
try ORTNative.check(
|
|
trainingApi.pointee.AddProperty?(handle, namePointer, OrtStringProperty, mutablePointer),
|
|
action: "AddProperty"
|
|
)
|
|
}
|
|
}
|
|
}
|
|
|
|
public func intProperty(name: String) throws -> Int64 {
|
|
guard let trainingApi = ORTNative.trainingApi else {
|
|
throw ORTError.unsupported("This ONNX Runtime build does not expose the training API")
|
|
}
|
|
|
|
return try ORTNative.withAllocator { allocator in
|
|
var propertyType = OrtPropertyType(rawValue: 0)
|
|
var propertyValue: UnsafeMutableRawPointer?
|
|
try name.withCString { pointer in
|
|
try ORTNative.check(
|
|
trainingApi.pointee.GetProperty?(handle, pointer, allocator, &propertyType, &propertyValue),
|
|
action: "GetProperty"
|
|
)
|
|
}
|
|
|
|
defer {
|
|
if let propertyValue {
|
|
_ = ORTNative.api.pointee.AllocatorFree?(allocator, propertyValue)
|
|
}
|
|
}
|
|
|
|
guard propertyType == OrtIntProperty, let propertyValue else {
|
|
throw ORTError.invalidArgument("Checkpoint property \(name) is not an Int64")
|
|
}
|
|
return propertyValue.assumingMemoryBound(to: Int64.self).pointee
|
|
}
|
|
}
|
|
|
|
public func floatProperty(name: String) throws -> Float {
|
|
guard let trainingApi = ORTNative.trainingApi else {
|
|
throw ORTError.unsupported("This ONNX Runtime build does not expose the training API")
|
|
}
|
|
|
|
return try ORTNative.withAllocator { allocator in
|
|
var propertyType = OrtPropertyType(rawValue: 0)
|
|
var propertyValue: UnsafeMutableRawPointer?
|
|
try name.withCString { pointer in
|
|
try ORTNative.check(
|
|
trainingApi.pointee.GetProperty?(handle, pointer, allocator, &propertyType, &propertyValue),
|
|
action: "GetProperty"
|
|
)
|
|
}
|
|
|
|
defer {
|
|
if let propertyValue {
|
|
_ = ORTNative.api.pointee.AllocatorFree?(allocator, propertyValue)
|
|
}
|
|
}
|
|
|
|
guard propertyType == OrtFloatProperty, let propertyValue else {
|
|
throw ORTError.invalidArgument("Checkpoint property \(name) is not a Float")
|
|
}
|
|
return propertyValue.assumingMemoryBound(to: Float.self).pointee
|
|
}
|
|
}
|
|
|
|
public func stringProperty(name: String) throws -> String? {
|
|
guard let trainingApi = ORTNative.trainingApi else {
|
|
throw ORTError.unsupported("This ONNX Runtime build does not expose the training API")
|
|
}
|
|
|
|
return try ORTNative.withAllocator { allocator in
|
|
var propertyType = OrtPropertyType(rawValue: 0)
|
|
var propertyValue: UnsafeMutableRawPointer?
|
|
try name.withCString { pointer in
|
|
try ORTNative.check(
|
|
trainingApi.pointee.GetProperty?(handle, pointer, allocator, &propertyType, &propertyValue),
|
|
action: "GetProperty"
|
|
)
|
|
}
|
|
|
|
guard propertyType == OrtStringProperty else {
|
|
if let propertyValue {
|
|
_ = ORTNative.api.pointee.AllocatorFree?(allocator, propertyValue)
|
|
}
|
|
throw ORTError.invalidArgument("Checkpoint property \(name) is not a String")
|
|
}
|
|
|
|
let string = propertyValue.map { String(cString: $0.assumingMemoryBound(to: CChar.self)) }
|
|
if let propertyValue {
|
|
_ = ORTNative.api.pointee.AllocatorFree?(allocator, propertyValue)
|
|
}
|
|
return string
|
|
}
|
|
}
|
|
}
|
|
|
|
public final class ORTTrainingSession {
|
|
private let handle: OpaquePointer
|
|
private let checkpoint: ORTCheckpoint
|
|
|
|
public init(
|
|
env: ORTEnv,
|
|
sessionOptions: ORTSessionOptions? = nil,
|
|
checkpoint: ORTCheckpoint,
|
|
trainModelPath: String,
|
|
evalModelPath: String? = nil,
|
|
optimizerModelPath: String? = nil
|
|
) throws {
|
|
guard let trainingApi = ORTNative.trainingApi else {
|
|
throw ORTError.unsupported("This ONNX Runtime build does not expose the training API")
|
|
}
|
|
|
|
self.checkpoint = checkpoint
|
|
|
|
let createdOptions = sessionOptions == nil ? try ORTSessionOptions() : nil
|
|
let optionsHandle = sessionOptions?.handle ?? createdOptions!.handle
|
|
var trainingSession: OpaquePointer?
|
|
|
|
try trainModelPath.withCString { trainPointer in
|
|
try evalModelPath?.withCString { evalPointer in
|
|
try optimizerModelPath?.withCString { optimizerPointer in
|
|
try ORTNative.check(
|
|
trainingApi.pointee.CreateTrainingSession?(
|
|
env.handle,
|
|
optionsHandle,
|
|
checkpoint.handle,
|
|
trainPointer,
|
|
evalPointer,
|
|
optimizerPointer,
|
|
&trainingSession
|
|
),
|
|
action: "CreateTrainingSession"
|
|
)
|
|
}
|
|
} ?? {
|
|
try optimizerModelPath?.withCString { optimizerPointer in
|
|
try ORTNative.check(
|
|
trainingApi.pointee.CreateTrainingSession?(
|
|
env.handle,
|
|
optionsHandle,
|
|
checkpoint.handle,
|
|
trainPointer,
|
|
nil,
|
|
optimizerPointer,
|
|
&trainingSession
|
|
),
|
|
action: "CreateTrainingSession"
|
|
)
|
|
} ?? {
|
|
try ORTNative.check(
|
|
trainingApi.pointee.CreateTrainingSession?(
|
|
env.handle,
|
|
optionsHandle,
|
|
checkpoint.handle,
|
|
trainPointer,
|
|
nil,
|
|
nil,
|
|
&trainingSession
|
|
),
|
|
action: "CreateTrainingSession"
|
|
)
|
|
}()
|
|
}()
|
|
}
|
|
|
|
guard let trainingSession else {
|
|
throw ORTError.runtimeFailure(action: "CreateTrainingSession", message: "Training session was not returned")
|
|
}
|
|
handle = trainingSession
|
|
}
|
|
|
|
deinit {
|
|
ORTNative.trainingApi?.pointee.ReleaseTrainingSession?(handle)
|
|
}
|
|
|
|
public func trainStep(inputValues: [ORTValue]) throws -> [ORTValue] {
|
|
try runStep(
|
|
countAction: "TrainingSessionGetTrainingModelOutputCount",
|
|
countCall: { out in ORTNative.trainingApi!.pointee.TrainingSessionGetTrainingModelOutputCount?(handle, out) },
|
|
stepAction: "TrainStep",
|
|
stepCall: { outputCount, outputBuffer in
|
|
let inputs = inputValues.map(\.handle)
|
|
let inputPointers = inputs.map(Optional.init)
|
|
return try inputPointers.withUnsafeBufferPointer { inputBuffer in
|
|
try ORTNative.check(
|
|
ORTNative.trainingApi!.pointee.TrainStep?(
|
|
handle,
|
|
nil,
|
|
inputs.count,
|
|
inputBuffer.baseAddress,
|
|
outputCount,
|
|
outputBuffer
|
|
),
|
|
action: "TrainStep"
|
|
)
|
|
}
|
|
}
|
|
)
|
|
}
|
|
|
|
public func evalStep(inputValues: [ORTValue]) throws -> [ORTValue] {
|
|
try runStep(
|
|
countAction: "TrainingSessionGetEvalModelOutputCount",
|
|
countCall: { out in ORTNative.trainingApi!.pointee.TrainingSessionGetEvalModelOutputCount?(handle, out) },
|
|
stepAction: "EvalStep",
|
|
stepCall: { outputCount, outputBuffer in
|
|
let inputs = inputValues.map(\.handle)
|
|
let inputPointers = inputs.map(Optional.init)
|
|
return try inputPointers.withUnsafeBufferPointer { inputBuffer in
|
|
try ORTNative.check(
|
|
ORTNative.trainingApi!.pointee.EvalStep?(
|
|
handle,
|
|
nil,
|
|
inputs.count,
|
|
inputBuffer.baseAddress,
|
|
outputCount,
|
|
outputBuffer
|
|
),
|
|
action: "EvalStep"
|
|
)
|
|
}
|
|
}
|
|
)
|
|
}
|
|
|
|
public func lazyResetGrad() throws {
|
|
guard let trainingApi = ORTNative.trainingApi else {
|
|
throw ORTError.unsupported("This ONNX Runtime build does not expose the training API")
|
|
}
|
|
try ORTNative.check(trainingApi.pointee.LazyResetGrad?(handle), action: "LazyResetGrad")
|
|
}
|
|
|
|
public func optimizerStep() throws {
|
|
guard let trainingApi = ORTNative.trainingApi else {
|
|
throw ORTError.unsupported("This ONNX Runtime build does not expose the training API")
|
|
}
|
|
try ORTNative.check(trainingApi.pointee.OptimizerStep?(handle, nil), action: "OptimizerStep")
|
|
}
|
|
|
|
public func trainInputNames() throws -> [String] {
|
|
try trainingNames(
|
|
countAction: "TrainingSessionGetTrainingModelInputCount",
|
|
countCall: { out in ORTNative.trainingApi!.pointee.TrainingSessionGetTrainingModelInputCount?(handle, out) },
|
|
nameAction: "TrainingSessionGetTrainingModelInputName",
|
|
nameCall: { index, allocator, out in
|
|
ORTNative.trainingApi!.pointee.TrainingSessionGetTrainingModelInputName?(handle, index, allocator, out)
|
|
}
|
|
)
|
|
}
|
|
|
|
public func evalInputNames() throws -> [String] {
|
|
try trainingNames(
|
|
countAction: "TrainingSessionGetEvalModelInputCount",
|
|
countCall: { out in ORTNative.trainingApi!.pointee.TrainingSessionGetEvalModelInputCount?(handle, out) },
|
|
nameAction: "TrainingSessionGetEvalModelInputName",
|
|
nameCall: { index, allocator, out in
|
|
ORTNative.trainingApi!.pointee.TrainingSessionGetEvalModelInputName?(handle, index, allocator, out)
|
|
}
|
|
)
|
|
}
|
|
|
|
public func trainOutputNames() throws -> [String] {
|
|
try trainingNames(
|
|
countAction: "TrainingSessionGetTrainingModelOutputCount",
|
|
countCall: { out in ORTNative.trainingApi!.pointee.TrainingSessionGetTrainingModelOutputCount?(handle, out) },
|
|
nameAction: "TrainingSessionGetTrainingModelOutputName",
|
|
nameCall: { index, allocator, out in
|
|
ORTNative.trainingApi!.pointee.TrainingSessionGetTrainingModelOutputName?(handle, index, allocator, out)
|
|
}
|
|
)
|
|
}
|
|
|
|
public func evalOutputNames() throws -> [String] {
|
|
try trainingNames(
|
|
countAction: "TrainingSessionGetEvalModelOutputCount",
|
|
countCall: { out in ORTNative.trainingApi!.pointee.TrainingSessionGetEvalModelOutputCount?(handle, out) },
|
|
nameAction: "TrainingSessionGetEvalModelOutputName",
|
|
nameCall: { index, allocator, out in
|
|
ORTNative.trainingApi!.pointee.TrainingSessionGetEvalModelOutputName?(handle, index, allocator, out)
|
|
}
|
|
)
|
|
}
|
|
|
|
public func registerLinearLRScheduler(
|
|
warmupStepCount: Int64,
|
|
totalStepCount: Int64,
|
|
initialLr: Float
|
|
) throws {
|
|
guard let trainingApi = ORTNative.trainingApi else {
|
|
throw ORTError.unsupported("This ONNX Runtime build does not expose the training API")
|
|
}
|
|
try ORTNative.check(
|
|
trainingApi.pointee.RegisterLinearLRScheduler?(handle, warmupStepCount, totalStepCount, initialLr),
|
|
action: "RegisterLinearLRScheduler"
|
|
)
|
|
}
|
|
|
|
public func schedulerStep() throws {
|
|
guard let trainingApi = ORTNative.trainingApi else {
|
|
throw ORTError.unsupported("This ONNX Runtime build does not expose the training API")
|
|
}
|
|
try ORTNative.check(trainingApi.pointee.SchedulerStep?(handle), action: "SchedulerStep")
|
|
}
|
|
|
|
public func learningRate() throws -> Float {
|
|
guard let trainingApi = ORTNative.trainingApi else {
|
|
throw ORTError.unsupported("This ONNX Runtime build does not expose the training API")
|
|
}
|
|
var learningRate: Float = 0
|
|
try ORTNative.check(trainingApi.pointee.GetLearningRate?(handle, &learningRate), action: "GetLearningRate")
|
|
return learningRate
|
|
}
|
|
|
|
public func setLearningRate(_ lr: Float) throws {
|
|
guard let trainingApi = ORTNative.trainingApi else {
|
|
throw ORTError.unsupported("This ONNX Runtime build does not expose the training API")
|
|
}
|
|
try ORTNative.check(trainingApi.pointee.SetLearningRate?(handle, lr), action: "SetLearningRate")
|
|
}
|
|
|
|
public func fromBuffer(_ buffer: ORTValue) throws {
|
|
guard let trainingApi = ORTNative.trainingApi else {
|
|
throw ORTError.unsupported("This ONNX Runtime build does not expose the training API")
|
|
}
|
|
try ORTNative.check(trainingApi.pointee.CopyBufferToParameters?(handle, buffer.handle, false), action: "CopyBufferToParameters")
|
|
}
|
|
|
|
public func toBuffer(onlyTrainable: Bool) throws -> ORTValue {
|
|
guard let trainingApi = ORTNative.trainingApi else {
|
|
throw ORTError.unsupported("This ONNX Runtime build does not expose the training API")
|
|
}
|
|
|
|
var parameterCount = 0
|
|
try ORTNative.check(
|
|
trainingApi.pointee.GetParametersSize?(handle, ¶meterCount, onlyTrainable),
|
|
action: "GetParametersSize"
|
|
)
|
|
|
|
let byteCount = parameterCount * MemoryLayout<Float>.stride
|
|
let buffer = NSMutableData(length: byteCount) ?? NSMutableData()
|
|
let value = try ORTValue(tensorData: buffer, elementType: .float, shape: [parameterCount])
|
|
try ORTNative.check(
|
|
trainingApi.pointee.CopyParametersToBuffer?(handle, value.handle, onlyTrainable),
|
|
action: "CopyParametersToBuffer"
|
|
)
|
|
return value
|
|
}
|
|
|
|
public func exportModelForInference(
|
|
outputPath: String,
|
|
graphOutputNames: [String]
|
|
) throws {
|
|
guard let trainingApi = ORTNative.trainingApi else {
|
|
throw ORTError.unsupported("This ONNX Runtime build does not expose the training API")
|
|
}
|
|
|
|
try outputPath.withCString { pathPointer in
|
|
try ORTNative.withCStringArray(graphOutputNames) { namePointers in
|
|
try namePointers.withUnsafeBufferPointer { nameBuffer in
|
|
try ORTNative.check(
|
|
trainingApi.pointee.ExportModelForInferencing?(
|
|
handle,
|
|
pathPointer,
|
|
graphOutputNames.count,
|
|
nameBuffer.baseAddress
|
|
),
|
|
action: "ExportModelForInferencing"
|
|
)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
private func trainingNames(
|
|
countAction: String,
|
|
countCall: (_ out: UnsafeMutablePointer<Int>) -> OpaquePointer?,
|
|
nameAction: String,
|
|
nameCall: (
|
|
_ index: Int,
|
|
_ allocator: UnsafeMutablePointer<OrtAllocator>,
|
|
_ out: UnsafeMutablePointer<UnsafeMutablePointer<CChar>?>
|
|
) -> OpaquePointer?
|
|
) throws -> [String] {
|
|
var count = 0
|
|
try ORTNative.check(countCall(&count), action: countAction)
|
|
|
|
return try ORTNative.withAllocator { allocator in
|
|
var names: [String] = []
|
|
names.reserveCapacity(count)
|
|
for index in 0..<count {
|
|
var namePointer: UnsafeMutablePointer<CChar>?
|
|
try ORTNative.check(nameCall(index, allocator, &namePointer), action: nameAction)
|
|
names.append(try ORTNative.copyAllocatedString(namePointer, allocator: allocator))
|
|
}
|
|
return names
|
|
}
|
|
}
|
|
|
|
private func runStep(
|
|
countAction: String,
|
|
countCall: (_ out: UnsafeMutablePointer<Int>) -> OpaquePointer?,
|
|
stepAction: String,
|
|
stepCall: (_ outputCount: Int, _ outputBuffer: UnsafeMutablePointer<OpaquePointer?>?) throws -> Void
|
|
) throws -> [ORTValue] {
|
|
var outputCount = 0
|
|
try ORTNative.check(countCall(&outputCount), action: countAction)
|
|
|
|
var outputs = Array<OpaquePointer?>(repeating: nil, count: outputCount)
|
|
try outputs.withUnsafeMutableBufferPointer { buffer in
|
|
try stepCall(outputCount, buffer.baseAddress)
|
|
}
|
|
|
|
return try outputs.enumerated().map { index, handle in
|
|
guard let handle else {
|
|
throw ORTError.runtimeFailure(action: stepAction, message: "Output \(index) was not returned")
|
|
}
|
|
return ORTValue(adopting: handle)
|
|
}
|
|
}
|
|
}
|