jon.recoil.org

Tensor Addition

This notebook demonstrates basic ONNX Runtime inference in OCaml: creating tensors, loading a model, and running addition.

Setup

Load onnxrt, Note (FRP), and the widget library:

#require "onnxrt";; #require "note";; #require "js_top_worker-widget";;

Load the ONNX Runtime JavaScript library into the worker:

let () = Js_of_ocaml.Js.Unsafe.meth_call Js_of_ocaml.Js.Unsafe.global "importScripts" [| Js_of_ocaml.Js.Unsafe.inject (Js_of_ocaml.Js.string "https://cdn.jsdelivr.net/npm/onnxruntime-web@1.21.0/dist/ort.min.js") |] let () = let open Js_of_ocaml in let ort = Js.Unsafe.get Js.Unsafe.global (Js.string "ort") in let env = Js.Unsafe.get ort (Js.string "env") in let wasm = Js.Unsafe.get env (Js.string "wasm") in Js.Unsafe.set wasm (Js.string "wasmPaths") (Js.string "https://cdn.jsdelivr.net/npm/onnxruntime-web@1.21.0/dist/") let () = print_endline "ort.js loaded"

Create Tensors

Create two Float32 tensors A = [1, 2, 3] and B = [4, 5, 6]. Edit the values and re-run to try different inputs:

open Onnxrt let a_data = Bigarray.Array1.of_array Bigarray.float32 Bigarray.c_layout [| 1.0; 2.0; 3.0 |] let b_data = Bigarray.Array1.of_array Bigarray.float32 Bigarray.c_layout [| 4.0; 5.0; 6.0 |] let a = Tensor.of_bigarray1 Dtype.Float32 a_data ~dims:[| 3 |] let b = Tensor.of_bigarray1 Dtype.Float32 b_data ~dims:[| 3 |] let () = print_endline "Tensors A and B created"

Run Inference

Load the add.onnx model and compute C = A + B. A status widget updates reactively as the model loads and inference completes:

let status_e, send_status = Note.E.create () let status = Note.S.hold "Loading model..." status_e let status_view msg = let open Widget.View in Element { tag = "div"; attrs = [ Style ("padding", "0.75em 1em"); Style ("border-radius", "6px"); Style ("font-family", "monospace"); Style ("background", "#f0f4f8"); ]; children = [Text msg] } let () = Widget.display ~id:"result" ~handlers:[] (status_view "Loading model...") let _logr = Note.S.log (Note.S.map status_view status) (Widget.update ~id:"result") let () = Note.Logr.hold _logr let () = Lwt.async (fun () -> let open Lwt.Syntax in let* session = Session.create "add.onnx" () in send_status "Running inference..."; let* outputs = Session.run session [("A", a); ("B", b)] in let c = List.assoc "C" outputs in let c_data = Tensor.to_bigarray1_exn Dtype.Float32 c in let result = Printf.sprintf "C = [%g, %g, %g]" (Bigarray.Array1.get c_data 0) (Bigarray.Array1.get c_data 1) (Bigarray.Array1.get c_data 2) in Tensor.dispose a; Tensor.dispose b; Tensor.dispose c; let* () = Session.release session in send_status result; Lwt.return_unit)

You should see the widget above update to C = [5, 7, 9] — the element-wise sum of A and B.