Machine Learning Pipelines

Mix.install([
  {:ex_zarr, "~> 1.2"},
  # {:ex_zarr, path: Path.join(__DIR__, "../../..")},
  {:nx, "~> 0.7"}
])

Streaming tensors from chunks

alias ExZarr.Array
alias ExZarr.Gallery.Pack

{:ok, array} =
  Array.create(
    shape: {128, 8},
    chunks: {32, 8},
    dtype: :float32,
    compressor: :zlib,
    storage: :memory
  )

vals = for i <- 0..(128 * 8 - 1), do: i * 0.01
:ok = Array.set_slice(array, Pack.pack(vals, :float32), start: {0, 0}, stop: {128, 8})

batches =
  array
  |> Array.stream_chunks(concurrency: 2)
  |> Stream.map(fn {_index, data} -> Nx.from_binary(data, {:f, 32}) end)
  |> Stream.chunk_every(2)
  |> Enum.to_list()

%{batch_count: length(batches), first_batch_shapes: Enum.map(hd(batches), &Nx.shape/1)}

DataLoader (sample-level batches)

# Requires Nx; shuffled batches for training-style loops
array
|> ExZarr.Nx.DataLoader.shuffled_batch_stream(16, seed: 42)
|> Enum.take(2)
|> Enum.map(fn
  {:ok, batch} -> Nx.shape(batch)
  other -> other
end)

Broadway

For production prep, see docs/livebooks/broadway_pipeline.livemd and ExZarr.Broadway.start_chunk_pipeline/3.