Flux.jl: Data loading & preprocessing pipeline feature

[ageron]

The DataLoader is nice, but if I understand correctly it requires the dataset to fit in memory. For large datasets that don't fit in memory, it would be nice to have an easy way to load & preprocess the data efficiently, similar to TensorFlow's tf.data API. Maybe something like this exists already?

If not, perhaps one option would be to provide custom transducers to make it possible to write things like:

data = (csv_file_paths |> Shuffle(length(csv_file_paths)) |> Interleave(CSV.File; threads=4)
             |> Map(preprocess_sample) |> Shuffle(100_000) |> Batch(32) |> Prefetch(1))

This would load records from multiple files (in random file order), pick 4 randomly, interleave their records, preprocess every record, shuffle records using a 100,000 element buffer, and batch the records with batch size 32, and prefetch 1 batch (so the CPU can prepare the next batch while the GPU is working on the previous batch). Then the data could be used for training.

[pxl-th]

Right now it is possible to use DataLoader with datasets that do not fit into memory, but feels somewhat hacky... or not.
Let's say we have dataset of images, which we want to load on demand.
Then we can define a custom struct subtyping AbstractArray.

struct Dataset{T, N} <: AbstractArray{T, N}
    frame_template::FormatExpr
    targets::AbstractArray
end

Dataset{T}(frame_template, targets) where {T} =
    Dataset{T, 1}(frame_template, targets)

Then defining getindex function that describes how we load one item

function Base.getindex(d::Dataset{T}, i::Int) where {T}
    path = format(d.frame_template, i - 1)
    image = path |> FileIO.load |> Images.channelview .|> T
    image, d.targets[[i]]
end

and how we load mini-batch

function Base.getindex(d::Dataset{T}, ids::Array) where {T}
    x, y = d[ids[1]]
    xs_last_dim = ntuple(i -> Colon(), ndims(x))
    ys_last_dim = ntuple(i -> Colon(), ndims(y))

    xs = Array{T}(undef, size(x)..., length(ids))
    ys = Array{T}(undef, size(y)..., length(ids))

    xs[xs_last_dim..., 1] .= x
    ys[ys_last_dim..., 1] .= y

    for (i, id) in enumerate(ids[2:end])
        x, y = d[id]
        xs[xs_last_dim..., i + 1] .= x
        ys[ys_last_dim..., i + 1] .= y
    end
    xs, ys
end

And some helper functions

Base.IndexStyle(::Type{Dataset}) = IndexLinear()
Base.size(d::Dataset) = (length(d.targets),)
Base.length(d::Dataset) = length(d.targets)

This in some sense mimicks PyTorch's Dataset
and allows to use DataLoader with datasets that do not fit into memory

frame_template = FormatExpr(raw".\frames\frame-{:d}.jpg")  # Template path for images.
targets = load_from_txt(raw".\speed.txt")  # Array of targets

dataset = Dataset{Float32}(frame_template, targets)
loader = Flux.Data.DataLoader(dataset, batchsize=4, shuffle=true)
println("Loader length: $(length(loader))")
for (i, (x, y)) in enumerate(loader)
    i == 10 && break
    println("$i: $(size(x)) $(size(y))")
end

Output of the data from my example:

Loader length: 240
1: (3, 160, 320, 4) (1, 4)
2: (3, 160, 320, 4) (1, 4)
3: (3, 160, 320, 4) (1, 4)
4: (3, 160, 320, 4) (1, 4)
5: (3, 160, 320, 4) (1, 4)
6: (3, 160, 320, 4) (1, 4)
7: (3, 160, 320, 4) (1, 4)
8: (3, 160, 320, 4) (1, 4)
9: (3, 160, 320, 4) (1, 4)

[CarloLucibello]

For the time being, we can just document the interface that a "Dataset" should expose in order to be compatible with the DataLoader. @pxl-th a PR in this direction would be very welcome.

In the longer run, we should definitely consider reimplementing the DataLoader on top of transducers. Transducers are great and come fully packed with features, as @ageron showed.

[ageron]

Thanks for your detailed answer @pxl-th . I'm not sure whether my code example really makes sense, but it's the kind of API I would imagine, largely inspired from TF's tf.data API, with a transducers twist. I'm happy to help if you want.

[darleybarreto]

I wonder if something more idiomatic could be done, like:

# I can call custom now and it will return three objects
@dataset (:train) image,target1,target2 function custom(path_arrays,idx)
    image = # load from the path_arrays[idx] + ...
    target1 = # load from the path_arrays[idx] + ...
    target2 = # load from the path_arrays[idx] + ...
end

or

# add another method to dataset
function dataset(path_arrays,idx, :train)
    image = # load from the path_arrays[idx] + ...
    target1 = # load from the path_arrays[idx] + ...
    target2 = # load from the path_arrays[idx] + ...

    image,target1,target2 
end

And another type called DataSet could be added having an inner field (_data or something like that), which could be the the array of paths (path_arrays) or the true data, depending on the user choice. And the data inside the DataLoader is a DataSet, which samples from the corresponding dataset in :train, :val, :test or any other custom symbol defining a step.

I am not a Julia expert, but I could help implementing it :smile: .