jon.recoil.org
A fork of mtelver's day10 project
1open Js_of_ocaml
2open Onnxrt
3
4let log msg =
5 let el = Dom_html.getElementById "output" in
6 let prev =
7 Js.to_string (Js.Opt.get el##.textContent (fun () -> Js.string ""))
8 in
9 el##.textContent := Js.some (Js.string (prev ^ msg ^ "\n"))
10
11let () =
12 Lwt.async @@ fun () ->
13 let open Lwt.Syntax in
14 log "Creating tensors...";
15 let a_data =
16 Bigarray.Array1.of_array Bigarray.float32 Bigarray.c_layout
17 [| 1.0; 2.0; 3.0 |]
18 in
19 let b_data =
20 Bigarray.Array1.of_array Bigarray.float32 Bigarray.c_layout
21 [| 4.0; 5.0; 6.0 |]
22 in
23 let a = Tensor.of_bigarray1 Dtype.Float32 a_data ~dims:[| 3 |] in
24 let b = Tensor.of_bigarray1 Dtype.Float32 b_data ~dims:[| 3 |] in
25 log "Loading model...";
26 let* session = Session.create "add.onnx" () in
27 log "Running inference...";
28 let* outputs = Session.run session [ ("A", a); ("B", b) ] in
29 let c = List.assoc "C" outputs in
30 let c_data = Tensor.to_bigarray1_exn Dtype.Float32 c in
31 log
32 (Printf.sprintf "Result C = [%g, %g, %g]" (Bigarray.Array1.get c_data 0)
33 (Bigarray.Array1.get c_data 1)
34 (Bigarray.Array1.get c_data 2));
35 Tensor.dispose a;
36 Tensor.dispose b;
37 Tensor.dispose c;
38 let* () = Session.release session in
39 log "Done!";
40 Lwt.return_unit