Machine Learning in Elixir: Chapter 7 train_pipeline example error

In page 149 of the book, there’s an example of the output of running the CatsAndDogs.pipeline function. The ouptut shows that the image tensors have a type and shape of f32[channels: 3][height: 96][width: 96] but those are channels-first tensors and the StbImage library returns channels-last tensors.

2 Likes

If you pattern match on the shape (and the type for good measure) of the %StbImage{…} your def pipeline(…) could read as:

  def pipeline(paths, batch_size, target_height, target_width) do
    paths
    |> Enum.shuffle()
    |> Task.async_stream(&parse_image/1)
    |> Stream.filter(fn
      # {:ok, {%StbImage{}, _}} -> true
      {:ok, {%StbImage{data: _, shape: {_, _, 3}, type: {:u, 8}}, _}} -> true
      _ -> false
    end)
    |> Stream.map(&to_tensors(&1, target_height, target_width))
    |> Stream.chunk_every(batch_size, batch_size, :discard)
    |> Stream.map(fn chunks ->
      {img_chunk, label_chunk} = Enum.unzip(chunks)
      {Nx.stack(img_chunk), Nx.stack(label_chunk)}
    end)
  end

The commented line is just to show what I changed. Cheers!