Class: OnnxRuby::Session

Inherits:
Object
  • Object
show all
Defined in:
lib/onnx_ruby/session.rb

Constant Summary collapse

VALID_PROVIDERS = 
%i[cpu coreml cuda tensorrt].freeze

Instance Method Summary collapse

Constructor Details

#initialize(model_path, providers: [:cpu], inter_threads: nil, intra_threads: nil, log_level: :warning, optimization_level: :all, memory_pattern: true, cpu_mem_arena: true, execution_mode: :sequential) ⇒ Session

Returns a new instance of Session.

Raises:



7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
 
# File 'lib/onnx_ruby/session.rb', line 7

def initialize(model_path, providers: [:cpu], inter_threads: nil, intra_threads: nil,
               log_level: :warning, optimization_level: :all, memory_pattern: true,
               cpu_mem_arena: true, execution_mode: :sequential)
  model_path = File.expand_path(model_path)
  raise ModelError, "model file not found: #{model_path}" unless File.exist?(model_path)

  provider_strs = Array(providers).map do |p|
    p = p.to_sym
    raise Error, "unknown provider: #{p}. Valid: #{VALID_PROVIDERS.join(", ")}" unless VALID_PROVIDERS.include?(p)
    p.to_s
  end

  @session = Ext::SessionWrapper.new(
    model_path,
    log_level_to_int(log_level),
    intra_threads || 0,
    inter_threads || 0,
    optimization_level.to_s,
    memory_pattern,
    cpu_mem_arena,
    execution_mode.to_s,
    provider_strs
  )
end

Instance Method Details

#inputs ⇒ Object 



32
33
34
 
# File 'lib/onnx_ruby/session.rb', line 32

def inputs
  @session.input_info
end

#outputs ⇒ Object 



36
37
38
 
# File 'lib/onnx_ruby/session.rb', line 36

def outputs
  @session.output_info
end

#run(inputs, output_names: nil) ⇒ Object 



40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
 
# File 'lib/onnx_ruby/session.rb', line 40

def run(inputs, output_names: nil)
  input_values = inputs.map do |name, data|
    if data.is_a?(Tensor)
      { name: name, data: data.flat_data, shape: data.shape, dtype: data.dtype.to_s }
    else
      shape = infer_shape(data)
      flat = data.flatten
      expected_size = shape.reduce(1, :*)
      if flat.length != expected_size
        raise TensorError,
              "input '#{name}' data size #{flat.length} does not match shape #{shape} (expected #{expected_size})"
      end
      dtype = infer_dtype(flat)
      { name: name, data: flat, shape: shape, dtype: dtype }
    end
  end

  @session.run(input_values, output_names || [])
end