The Nx_io
module provides functions to load and save Nx tensors in various file formats, including image formats and NumPy formats.
Features
- Image I/O: Load and save images in PNG, JPEG, BMP, TGA, and GIF formats
- NumPy .npy format: Load and save single arrays in NumPy's native format
- NumPy .npz archives: Load and save multiple named arrays in compressed archives
- Runtime dtype detection: Handle arrays with types determined at runtime
- Type conversion utilities: Convert between different numeric types
Image Operations
Loading Images
Images can be loaded as uint8 tensors:
open Nx_io
(* Load as RGB: shape [|height; width; 3|] *)
let img = load_image "photo.png"
(* Load as grayscale: shape [|height; width|] *)
let gray = load_image ~grayscale:true "photo.png"
Saving Images
Save uint8 tensors as images:
(* Save RGB or grayscale based on shape *)
save_image img "output.png"
Supported shapes:
|height; width|
- Grayscale|height; width; 1|
- Grayscale with explicit channel|height; width; 3|
- RGB|height; width; 4|
- RGBA
Single Arrays (.npy)
The .npy format stores a single array with its dtype and shape information:
(* Load array with runtime-detected type *)
let P arr = load_npy "data.npy" in
(* arr : ('a, 'b) Nx.t *)
(* Convert to specific type *)
let float_arr = load_npy "data.npy" |> to_float32
(* Save array *)
save_npy my_array "output.npy"
Archives (.npz)
The .npz format stores multiple named arrays in a compressed archive:
(* Load entire archive *)
let archive = load_npz "bundle.npz" in
(* Access specific array *)
match Hashtbl.find_opt archive "weights" with
| Some (P arr) ->
let weights = to_float32 (P arr) in
(* use weights *)
| None -> failwith "weights not found"
(* Load single array directly *)
let P data = load_npz_member ~path:"bundle.npz" ~name:"data"
(* Save multiple arrays *)
save_npz [
("inputs", P input_array);
("labels", P label_array);
("weights", P weight_array)
] "model.npz"
Packed Arrays and Type Conversions
Since file formats store type information that's only known at runtime, loaded arrays are wrapped in the Nx_io.packed_nx
type:
type packed_nx = P : ('a, 'b) Nx.t -> packed_nx
Convert packed arrays to specific types using the provided functions:
let packed = load_npy "data.npy" in
let float32_array = to_float32 packed
let int32_array = to_int32 packed
let uint8_array = to_uint8 packed
(* etc. *)
Available conversions:
- Floating point:
to_float16
, to_float32
, to_float64
- Signed integers:
to_int8
, to_int16
, to_int32
, to_int64
- Unsigned integers:
to_uint8
, to_uint16
- Complex:
to_complex32
, to_complex64
Examples
Image Processing Pipeline
open Nx
open Nx_io
let process_image input_path output_path =
(* Load image *)
let img = load_image input_path in
(* Convert to float for processing *)
let img_float = Nx.astype float32 img in
(* Normalize to [0, 1] *)
let normalized = Nx.div_s img_float 255.0 in
(* Apply some processing *)
let processed =
normalized
|> Nx.mul_s 0.8 (* Reduce brightness *)
|> Nx.add_s 0.1 (* Add bias *)
in
(* Convert back to uint8 *)
let result =
processed
|> Nx.mul_s 255.0
|> Nx.clip ~min:0.0 ~max:255.0
|> Nx.astype uint8
in
(* Save result *)
save_image result output_path
Model Checkpoint Save/Load
let save_checkpoint ~path ~epoch ~model =
let weights = Model.get_weights model in
let optimizer_state = Model.get_optimizer_state model in
save_npz [
("epoch", P (Nx.scalar int32 epoch));
("weights", P weights);
("optimizer_state", P optimizer_state);
] path
let load_checkpoint path =
let archive = load_npz path in
let epoch =
match Hashtbl.find_opt archive "epoch" with
| Some p -> Nx.get_item [] (to_int32 p)
| None -> failwith "epoch not found"
in
let weights =
match Hashtbl.find_opt archive "weights" with
| Some p -> to_float32 p
| None -> failwith "weights not found"
in
(epoch, weights)
Error Handling
All I/O operations may raise Failure
exceptions:
- File not found or inaccessible
- Unsupported file format
- Invalid data format
- Incompatible array shapes (for
save_image
) - Missing archive members (for
load_npz_member
)
- Image loading/saving uses the stb_image libraries (header-only C libraries)
- NumPy format I/O is implemented in pure OCaml
- Large .npz archives are loaded entirely into memory
- For very large datasets, consider loading arrays individually with
load_npz_member
See Also
Nx
: Core tensor operationsNx_io
: Complete API reference