module OnnxRuntime class InferenceSession attr_reader :inputs, :outputs def initialize(path_or_bytes, enable_cpu_mem_arena: true, enable_mem_pattern: true, enable_profiling: false, execution_mode: nil, free_dimension_overrides_by_denotation: nil, free_dimension_overrides_by_name: nil, graph_optimization_level: nil, inter_op_num_threads: nil, intra_op_num_threads: nil, log_severity_level: nil, log_verbosity_level: nil, logid: nil, optimized_model_filepath: nil, profile_file_prefix: nil, session_config_entries: nil, providers: []) # session options session_options = ::FFI::MemoryPointer.new(:pointer) check_status api[:CreateSessionOptions].call(session_options) if enable_cpu_mem_arena check_status api[:EnableCpuMemArena].call(session_options.read_pointer) else check_status api[:DisableCpuMemArena].call(session_options.read_pointer) end if enable_mem_pattern check_status api[:EnableMemPattern].call(session_options.read_pointer) else check_status api[:DisableMemPattern].call(session_options.read_pointer) end if enable_profiling check_status api[:EnableProfiling].call(session_options.read_pointer, ort_string(profile_file_prefix || "onnxruntime_profile_")) else check_status api[:DisableProfiling].call(session_options.read_pointer) end if execution_mode execution_modes = {sequential: 0, parallel: 1} mode = execution_modes[execution_mode] raise ArgumentError, "Invalid execution mode" unless mode check_status api[:SetSessionExecutionMode].call(session_options.read_pointer, mode) end if free_dimension_overrides_by_denotation free_dimension_overrides_by_denotation.each do |k, v| check_status api[:AddFreeDimensionOverride].call(session_options.read_pointer, k.to_s, v) end end if free_dimension_overrides_by_name free_dimension_overrides_by_name.each do |k, v| check_status api[:AddFreeDimensionOverrideByName].call(session_options.read_pointer, k.to_s, v) end end if graph_optimization_level optimization_levels = {none: 0, basic: 1, extended: 2, all: 99} level = optimization_levels[graph_optimization_level] raise ArgumentError, "Invalid graph optimization level" unless level check_status api[:SetSessionGraphOptimizationLevel].call(session_options.read_pointer, level) end check_status api[:SetInterOpNumThreads].call(session_options.read_pointer, inter_op_num_threads) if inter_op_num_threads check_status api[:SetIntraOpNumThreads].call(session_options.read_pointer, intra_op_num_threads) if intra_op_num_threads check_status api[:SetSessionLogSeverityLevel].call(session_options.read_pointer, log_severity_level) if log_severity_level check_status api[:SetSessionLogVerbosityLevel].call(session_options.read_pointer, log_verbosity_level) if log_verbosity_level check_status api[:SetSessionLogId].call(session_options.read_pointer, logid) if logid check_status api[:SetOptimizedModelFilePath].call(session_options.read_pointer, ort_string(optimized_model_filepath)) if optimized_model_filepath if session_config_entries session_config_entries.each do |k, v| check_status api[:AddSessionConfigEntry].call(session_options.read_pointer, k.to_s, v.to_s) end end providers.each do |provider| unless self.providers.include?(provider) warn "Provider not available: #{provider}" next end case provider when "CUDAExecutionProvider" cuda_options = ::FFI::MemoryPointer.new(:pointer) check_status api[:CreateCUDAProviderOptions].call(cuda_options) check_status api[:SessionOptionsAppendExecutionProvider_CUDA_V2].call(session_options.read_pointer, cuda_options.read_pointer) release :CUDAProviderOptions, cuda_options when "CoreMLExecutionProvider" unless FFI.respond_to?(:OrtSessionOptionsAppendExecutionProvider_CoreML) raise ArgumentError, "Provider not available: #{provider}" end coreml_flags = 0 check_status FFI.OrtSessionOptionsAppendExecutionProvider_CoreML(session_options.read_pointer, coreml_flags) when "CPUExecutionProvider" break else raise ArgumentError, "Provider not supported: #{provider}" end end @session = load_session(path_or_bytes, session_options) ObjectSpace.define_finalizer(@session, self.class.finalize(read_pointer.to_i)) @allocator = load_allocator @inputs = load_inputs @outputs = load_outputs ensure release :SessionOptions, session_options end # TODO support logid def run(output_names, input_feed, log_severity_level: nil, log_verbosity_level: nil, logid: nil, terminate: nil, output_type: :ruby) # pointer references refs = [] input_tensor = create_input_tensor(input_feed, refs) output_names ||= @outputs.map { |v| v[:name] } output_tensor = ::FFI::MemoryPointer.new(:pointer, outputs.size) input_node_names = create_node_names(input_feed.keys.map(&:to_s), refs) output_node_names = create_node_names(output_names.map(&:to_s), refs) # run options run_options = ::FFI::MemoryPointer.new(:pointer) check_status api[:CreateRunOptions].call(run_options) check_status api[:RunOptionsSetRunLogSeverityLevel].call(run_options.read_pointer, log_severity_level) if log_severity_level check_status api[:RunOptionsSetRunLogVerbosityLevel].call(run_options.read_pointer, log_verbosity_level) if log_verbosity_level check_status api[:RunOptionsSetRunTag].call(run_options.read_pointer, logid) if logid check_status api[:RunOptionsSetTerminate].call(run_options.read_pointer) if terminate check_status api[:Run].call(read_pointer, run_options.read_pointer, input_node_names, input_tensor, input_feed.size, output_node_names, output_names.size, output_tensor) output_names.size.times.map do |i| create_from_onnx_value(output_tensor[i].read_pointer, output_type) end ensure release :RunOptions, run_options if input_tensor input_feed.size.times do |i| release :Value, input_tensor[i] end end # output values released in create_from_onnx_value end def modelmeta keys = ::FFI::MemoryPointer.new(:pointer) num_keys = ::FFI::MemoryPointer.new(:int64_t) description = ::FFI::MemoryPointer.new(:string) domain = ::FFI::MemoryPointer.new(:string) graph_name = ::FFI::MemoryPointer.new(:string) graph_description = ::FFI::MemoryPointer.new(:string) producer_name = ::FFI::MemoryPointer.new(:string) version = ::FFI::MemoryPointer.new(:int64_t) metadata = ::FFI::MemoryPointer.new(:pointer) check_status api[:SessionGetModelMetadata].call(read_pointer, metadata) custom_metadata_map = {} check_status api[:ModelMetadataGetCustomMetadataMapKeys].call(metadata.read_pointer, @allocator.read_pointer, keys, num_keys) num_keys.read(:int64_t).times do |i| key_ptr = keys.read_pointer[i * ::FFI::Pointer.size] key = key_ptr.read_pointer.read_string value = ::FFI::MemoryPointer.new(:string) check_status api[:ModelMetadataLookupCustomMetadataMap].call(metadata.read_pointer, @allocator.read_pointer, key, value) custom_metadata_map[key] = value.read_pointer.read_string allocator_free key_ptr allocator_free value end allocator_free keys check_status api[:ModelMetadataGetDescription].call(metadata.read_pointer, @allocator.read_pointer, description) check_status api[:ModelMetadataGetDomain].call(metadata.read_pointer, @allocator.read_pointer, domain) check_status api[:ModelMetadataGetGraphName].call(metadata.read_pointer, @allocator.read_pointer, graph_name) check_status api[:ModelMetadataGetGraphDescription].call(metadata.read_pointer, @allocator.read_pointer, graph_description) check_status api[:ModelMetadataGetProducerName].call(metadata.read_pointer, @allocator.read_pointer, producer_name) check_status api[:ModelMetadataGetVersion].call(metadata.read_pointer, version) { custom_metadata_map: custom_metadata_map, description: description.read_pointer.read_string, domain: domain.read_pointer.read_string, graph_name: graph_name.read_pointer.read_string, graph_description: graph_description.read_pointer.read_string, producer_name: producer_name.read_pointer.read_string, version: version.read(:int64_t) } ensure release :ModelMetadata, metadata allocator_free description allocator_free domain allocator_free graph_name allocator_free graph_description allocator_free producer_name end # return value has double underscore like Python def end_profiling out = ::FFI::MemoryPointer.new(:string) check_status api[:SessionEndProfiling].call(read_pointer, @allocator.read_pointer, out) out.read_pointer.read_string end # no way to set providers with C API yet # so we can return all available providers def providers out_ptr = ::FFI::MemoryPointer.new(:pointer) length_ptr = ::FFI::MemoryPointer.new(:int) check_status api[:GetAvailableProviders].call(out_ptr, length_ptr) length = length_ptr.read_int providers = [] length.times do |i| providers << out_ptr.read_pointer[i * ::FFI::Pointer.size].read_pointer.read_string end api[:ReleaseAvailableProviders].call(out_ptr.read_pointer, length) providers end private def load_session(path_or_bytes, session_options) session = ::FFI::MemoryPointer.new(:pointer) from_memory = if path_or_bytes.respond_to?(:read) path_or_bytes = path_or_bytes.read true else path_or_bytes = path_or_bytes.to_str false end if from_memory check_status api[:CreateSessionFromArray].call(env.read_pointer, path_or_bytes, path_or_bytes.bytesize, session_options.read_pointer, session) else check_status api[:CreateSession].call(env.read_pointer, ort_string(path_or_bytes), session_options.read_pointer, session) end session end def load_allocator allocator = ::FFI::MemoryPointer.new(:pointer) check_status api[:GetAllocatorWithDefaultOptions].call(allocator) allocator end def load_inputs inputs = [] num_input_nodes = ::FFI::MemoryPointer.new(:size_t) check_status api[:SessionGetInputCount].call(read_pointer, num_input_nodes) num_input_nodes.read(:size_t).times do |i| name_ptr = ::FFI::MemoryPointer.new(:string) check_status api[:SessionGetInputName].call(read_pointer, i, @allocator.read_pointer, name_ptr) # freed in node_info typeinfo = ::FFI::MemoryPointer.new(:pointer) check_status api[:SessionGetInputTypeInfo].call(read_pointer, i, typeinfo) inputs << {name: name_ptr.read_pointer.read_string}.merge(node_info(typeinfo)) allocator_free name_ptr end inputs end def load_outputs outputs = [] num_output_nodes = ::FFI::MemoryPointer.new(:size_t) check_status api[:SessionGetOutputCount].call(read_pointer, num_output_nodes) num_output_nodes.read(:size_t).times do |i| name_ptr = ::FFI::MemoryPointer.new(:string) check_status api[:SessionGetOutputName].call(read_pointer, i, @allocator.read_pointer, name_ptr) # freed in node_info typeinfo = ::FFI::MemoryPointer.new(:pointer) check_status api[:SessionGetOutputTypeInfo].call(read_pointer, i, typeinfo) outputs << {name: name_ptr.read_pointer.read_string}.merge(node_info(typeinfo)) allocator_free name_ptr end outputs end def create_input_tensor(input_feed, refs) allocator_info = ::FFI::MemoryPointer.new(:pointer) check_status api[:CreateCpuMemoryInfo].call(1, 0, allocator_info) input_tensor = ::FFI::MemoryPointer.new(:pointer, input_feed.size) input_feed.each_with_index do |(input_name, input), idx| # TODO support more types inp = @inputs.find { |i| i[:name] == input_name.to_s } raise Error, "Unknown input: #{input_name}" unless inp input = input.to_a unless input.is_a?(Array) || numo_array?(input) shape = input_shape(input) input_node_dims = ::FFI::MemoryPointer.new(:int64, shape.size) input_node_dims.write_array_of_int64(shape) if inp[:type] == "tensor(string)" type_enum = FFI::TensorElementDataType[:string] check_status api[:CreateTensorAsOrtValue].call(@allocator.read_pointer, input_node_dims, shape.size, type_enum, input_tensor[idx]) # keep reference to _str_ptrs until FillStringTensor call input_tensor_values, _str_ptrs = create_input_strings(input) check_status api[:FillStringTensor].call(input_tensor[idx].read_pointer, input_tensor_values, input_tensor_values.size / input_tensor_values.type_size) elsif (tensor_type = tensor_types[inp[:type]]) input_tensor_values = create_input_data(input, tensor_type) type_enum = FFI::TensorElementDataType[tensor_type] check_status api[:CreateTensorWithDataAsOrtValue].call(allocator_info.read_pointer, input_tensor_values, input_tensor_values.size, input_node_dims, shape.size, type_enum, input_tensor[idx]) refs << input_tensor_values else unsupported_type("input", inp[:type]) end end input_tensor ensure release :MemoryInfo, allocator_info end def input_shape(input) if numo_array?(input) input.shape else shape = [] s = input while s.is_a?(Array) shape << s.size s = s.first end shape end end def create_input_strings(input) str_ptrs = if numo_array?(input) input.size.times.map { |i| ::FFI::MemoryPointer.from_string(input[i]) } else input.flatten.map { |v| ::FFI::MemoryPointer.from_string(v) } end input_tensor_values = ::FFI::MemoryPointer.new(:pointer, str_ptrs.size) input_tensor_values.write_array_of_pointer(str_ptrs) [input_tensor_values, str_ptrs] end def create_input_data(input, tensor_type) if numo_array?(input) input.cast_to(numo_types[tensor_type]).to_binary else flat_input = input.flatten.to_a input_tensor_values = ::FFI::MemoryPointer.new(tensor_type, flat_input.size) if tensor_type == :bool input_tensor_values.write_array_of_uint8(flat_input.map { |v| v ? 1 : 0 }) else input_tensor_values.send("write_array_of_#{tensor_type}", flat_input) end input_tensor_values end end def create_node_names(names, refs) str_ptrs = names.map { |v| ::FFI::MemoryPointer.from_string(v) } refs << str_ptrs ptr = ::FFI::MemoryPointer.new(:pointer, names.size) ptr.write_array_of_pointer(str_ptrs) ptr end def create_from_onnx_value(out_ptr, output_type) out_type = ::FFI::MemoryPointer.new(:int) check_status api[:GetValueType].call(out_ptr, out_type) type = FFI::OnnxType[out_type.read_int] case type when :tensor typeinfo = ::FFI::MemoryPointer.new(:pointer) check_status api[:GetTensorTypeAndShape].call(out_ptr, typeinfo) type, shape = tensor_type_and_shape(typeinfo) tensor_data = ::FFI::MemoryPointer.new(:pointer) check_status api[:GetTensorMutableData].call(out_ptr, tensor_data) out_size = ::FFI::MemoryPointer.new(:size_t) check_status api[:GetTensorShapeElementCount].call(typeinfo.read_pointer, out_size) output_tensor_size = out_size.read(:size_t) release :TensorTypeAndShapeInfo, typeinfo # TODO support more types type = FFI::TensorElementDataType[type] case output_type when :numo case type when :string result = Numo::RObject.new(shape) result.allocate create_strings_from_onnx_value(out_ptr, output_tensor_size, result) else numo_type = numo_types[type] unsupported_type("element", type) unless numo_type numo_type.from_binary(tensor_data.read_pointer.read_bytes(output_tensor_size * numo_type::ELEMENT_BYTE_SIZE), shape) end when :ruby arr = case type when :float, :uint8, :int8, :uint16, :int16, :int32, :int64, :double, :uint32, :uint64 tensor_data.read_pointer.send("read_array_of_#{type}", output_tensor_size) when :bool tensor_data.read_pointer.read_array_of_uint8(output_tensor_size).map { |v| v == 1 } when :string create_strings_from_onnx_value(out_ptr, output_tensor_size, []) else unsupported_type("element", type) end Utils.reshape(arr, shape) else raise ArgumentError, "Invalid output type: #{output_type}" end when :sequence out = ::FFI::MemoryPointer.new(:size_t) check_status api[:GetValueCount].call(out_ptr, out) out.read(:size_t).times.map do |i| seq = ::FFI::MemoryPointer.new(:pointer) check_status api[:GetValue].call(out_ptr, i, @allocator.read_pointer, seq) create_from_onnx_value(seq.read_pointer, output_type) end when :map type_shape = ::FFI::MemoryPointer.new(:pointer) map_keys = ::FFI::MemoryPointer.new(:pointer) map_values = ::FFI::MemoryPointer.new(:pointer) elem_type = ::FFI::MemoryPointer.new(:int) check_status api[:GetValue].call(out_ptr, 0, @allocator.read_pointer, map_keys) check_status api[:GetValue].call(out_ptr, 1, @allocator.read_pointer, map_values) check_status api[:GetTensorTypeAndShape].call(map_keys.read_pointer, type_shape) check_status api[:GetTensorElementType].call(type_shape.read_pointer, elem_type) release :TensorTypeAndShapeInfo, type_shape # TODO support more types elem_type = FFI::TensorElementDataType[elem_type.read_int] case elem_type when :int64 ret = {} keys = create_from_onnx_value(map_keys.read_pointer, output_type) values = create_from_onnx_value(map_values.read_pointer, output_type) keys.zip(values).each do |k, v| ret[k] = v end ret else unsupported_type("element", elem_type) end else unsupported_type("ONNX", type) end ensure api[:ReleaseValue].call(out_ptr) unless out_ptr.null? end def create_strings_from_onnx_value(out_ptr, output_tensor_size, result) len = ::FFI::MemoryPointer.new(:size_t) check_status api[:GetStringTensorDataLength].call(out_ptr, len) s_len = len.read(:size_t) s = ::FFI::MemoryPointer.new(:uchar, s_len) offsets = ::FFI::MemoryPointer.new(:size_t, output_tensor_size) check_status api[:GetStringTensorContent].call(out_ptr, s, s_len, offsets, output_tensor_size) offsets = output_tensor_size.times.map { |i| offsets[i].read(:size_t) } offsets << s_len output_tensor_size.times do |i| result[i] = s.get_bytes(offsets[i], offsets[i + 1] - offsets[i]) end result end def read_pointer @session.read_pointer end def check_status(status) unless status.null? message = api[:GetErrorMessage].call(status).read_string api[:ReleaseStatus].call(status) raise Error, message end end def node_info(typeinfo) onnx_type = ::FFI::MemoryPointer.new(:int) check_status api[:GetOnnxTypeFromTypeInfo].call(typeinfo.read_pointer, onnx_type) type = FFI::OnnxType[onnx_type.read_int] case type when :tensor tensor_info = ::FFI::MemoryPointer.new(:pointer) # don't free tensor_info check_status api[:CastTypeInfoToTensorInfo].call(typeinfo.read_pointer, tensor_info) type, shape = tensor_type_and_shape(tensor_info) { type: "tensor(#{FFI::TensorElementDataType[type]})", shape: shape } when :sequence sequence_type_info = ::FFI::MemoryPointer.new(:pointer) check_status api[:CastTypeInfoToSequenceTypeInfo].call(typeinfo.read_pointer, sequence_type_info) nested_type_info = ::FFI::MemoryPointer.new(:pointer) check_status api[:GetSequenceElementType].call(sequence_type_info.read_pointer, nested_type_info) v = node_info(nested_type_info)[:type] { type: "seq(#{v})", shape: [] } when :map map_type_info = ::FFI::MemoryPointer.new(:pointer) check_status api[:CastTypeInfoToMapTypeInfo].call(typeinfo.read_pointer, map_type_info) # key key_type = ::FFI::MemoryPointer.new(:int) check_status api[:GetMapKeyType].call(map_type_info.read_pointer, key_type) k = FFI::TensorElementDataType[key_type.read_int] # value value_type_info = ::FFI::MemoryPointer.new(:pointer) check_status api[:GetMapValueType].call(map_type_info.read_pointer, value_type_info) v = node_info(value_type_info)[:type] { type: "map(#{k},#{v})", shape: [] } else unsupported_type("ONNX", type) end ensure release :TypeInfo, typeinfo end def tensor_type_and_shape(tensor_info) type = ::FFI::MemoryPointer.new(:int) check_status api[:GetTensorElementType].call(tensor_info.read_pointer, type) num_dims_ptr = ::FFI::MemoryPointer.new(:size_t) check_status api[:GetDimensionsCount].call(tensor_info.read_pointer, num_dims_ptr) num_dims = num_dims_ptr.read(:size_t) node_dims = ::FFI::MemoryPointer.new(:int64, num_dims) check_status api[:GetDimensions].call(tensor_info.read_pointer, node_dims, num_dims) dims = node_dims.read_array_of_int64(num_dims) symbolic_dims = ::FFI::MemoryPointer.new(:pointer, num_dims) check_status api[:GetSymbolicDimensions].call(tensor_info.read_pointer, symbolic_dims, num_dims) named_dims = num_dims.times.map { |i| symbolic_dims[i].read_pointer.read_string } dims = named_dims.zip(dims).map { |n, d| n.empty? ? d : n } [type.read_int, dims] end def unsupported_type(name, type) raise Error, "Unsupported #{name} type: #{type}" end def tensor_types @tensor_types ||= [:float, :uint8, :int8, :uint16, :int16, :int32, :int64, :bool, :double, :uint32, :uint64].map { |v| ["tensor(#{v})", v] }.to_h end def numo_array?(obj) defined?(Numo::NArray) && obj.is_a?(Numo::NArray) end def numo_types @numo_types ||= { float: Numo::SFloat, uint8: Numo::UInt8, int8: Numo::Int8, uint16: Numo::UInt16, int16: Numo::Int16, int32: Numo::Int32, int64: Numo::Int64, bool: Numo::UInt8, double: Numo::DFloat, uint32: Numo::UInt32, uint64: Numo::UInt64 } end def api self.class.api end def release(*args) self.class.release(*args) end def allocator_free(ptr) api[:AllocatorFree].call(@allocator.read_pointer, ptr.read_pointer) end def self.api FFI.api end def self.release(type, pointer) api[:"Release#{type}"].call(pointer.read_pointer) if pointer && !pointer.null? end def self.finalize(addr) # must use proc instead of stabby lambda proc { api[:ReleaseSession].call(::FFI::Pointer.new(:pointer, addr)) } end # wide string on Windows # char string on Linux # see ORTCHAR_T in onnxruntime_c_api.h def ort_string(str) if Gem.win_platform? max = str.size + 1 # for null byte dest = ::FFI::MemoryPointer.new(:wchar_t, max) ret = FFI::Libc.mbstowcs(dest, str, max) raise Error, "Expected mbstowcs to return #{str.size}, got #{ret}" if ret != str.size dest else str end end def env # use mutex for thread-safety Utils.mutex.synchronize do @@env ||= begin env = ::FFI::MemoryPointer.new(:pointer) check_status api[:CreateEnv].call(3, "Default", env) at_exit { release :Env, env } # disable telemetry # https://github.com/microsoft/onnxruntime/blob/master/docs/Privacy.md check_status api[:DisableTelemetryEvents].call(env) env end end end end end