Expose api for custom datasets

docs/make.jl

@@ -15,7 +15,8 @@ makedocs(modules=[Flux, NNlib],
                      "NNlib" => "models/nnlib.md"],
                   "Handling Data" =>
                     ["One-Hot Encoding" => "data/onehot.md",
-                     "DataLoader" => "data/dataloader.md"],
+                     "DataLoader" => "data/dataloader.md",
+                     "Custom Dataset" => "data/dataset.md"],
                   "Training Models" =>
                     ["Optimisers" => "training/optimisers.md",
                      "Training" => "training/training.md"],

docs/src/data/dataset.md

@@ -0,0 +1,34 @@
+# Custom Dataset
+
+In order to maintain compatibility for custom datasets with `DataLoader`
+you need to implement following methods:
+
+- `Flux.Data.nobs(::CustomDataset)` -- total number of items in `CustomDataset`;
+- `Flux.Data.getobs(::CustomDataset, ids)` -- how to retrieve items from dataset for given list of `ids`;
+- `Base.eltype(::DataLoader{CustomDataset})` -- type of the elements returned by dataset.
+
+Below is a dummy example of how to adapt custom dataset
+to make it compatible with `DataLoader`.
+
+```julia
+# For each index returns an array of zeros.
+struct CustomDataset{T, N}
+    element_size::Tuple
+    total::Int
+end
+
+Base.eltype(::DataLoader{CustomDataset{T, N}}) where {T, N} = Array{T, N}
+
+Flux.Data.nobs(d::CustomDataset) = d.total
+function Flux.Data.getobs(d::CustomDataset{T, N}, i)::Array{T, N} where {T, N}
+    zeros(T, d.element_size..., length(i))
+end
+```
+
+And now you can use `CustomDataset` with `DataLoader`:
+
+```julia
+dataset = CustomDataset{Float32, 4}((28, 28, 1), 16)
+loader = DataLoader(dataset, batchsize=4, shuffle=true)
+batches = collect(loader)
+```

src/data/dataloader.jl

@@ -67,8 +67,8 @@ Usage example:
 """
 function DataLoader(data; batchsize=1, shuffle=false, partial=true)
     batchsize > 0 || throw(ArgumentError("Need positive batchsize"))
-    
-    n = _nobs(data) 
+
+    n = nobs(data)
     if n < batchsize
         @warn "Number of observations less than batchsize, decreasing the batchsize to $n"
         batchsize = n
@@ -84,7 +84,7 @@ end
     end
     nexti = min(i + d.batchsize, d.nobs)
     ids = d.indices[i+1:nexti]
-    batch = _getobs(d.data, ids)
+    batch = getobs(d.data, ids)
     return (batch, nexti)
 end
 
@@ -93,18 +93,18 @@ function Base.length(d::DataLoader)
     d.partial ? ceil(Int,n) : floor(Int,n)
 end
 
-_nobs(data::AbstractArray) = size(data)[end]
+nobs(data::AbstractArray) = size(data)[end]
 
-function _nobs(data::Union{Tuple, NamedTuple})
+function nobs(data::Union{Tuple, NamedTuple})
     length(data) > 0 || throw(ArgumentError("Need at least one data input"))
-    n = _nobs(data[1])
-    if !all(x -> _nobs(x) == n, Base.tail(data))
+    n = nobs(data[1])
+    if !all(x -> nobs(x) == n, Base.tail(data))
         throw(DimensionMismatch("All data should contain same number of observations"))
     end
     return n
 end
 
-_getobs(data::AbstractArray, i) = data[ntuple(i -> Colon(), Val(ndims(data) - 1))..., i]
-_getobs(data::Union{Tuple, NamedTuple}, i) = map(Base.Fix2(_getobs, i), data)
+getobs(data::AbstractArray, i) = data[ntuple(i -> Colon(), Val(ndims(data) - 1))..., i]
+getobs(data::Union{Tuple, NamedTuple}, i) = map(Base.Fix2(getobs, i), data)
 
 Base.eltype(::DataLoader{D}) where D = D

test/data.jl

@@ -76,6 +76,37 @@
     @test norm(θ .- 1) < 1e-10
 end
 
+@testset "Dataset" begin
+    struct ZerosDataset{T, N}
+        element_size::Tuple
+        total::Int
+    end
+
+    Base.eltype(::DataLoader{ZerosDataset{T, N}}) where {T, N} = Array{T, N}
+
+    Flux.Data.nobs(d::ZerosDataset) = d.total
+    function Flux.Data.getobs(d::ZerosDataset{T, N}, i)::Array{T, N} where {T, N}
+        zeros(T, d.element_size..., length(i))
+    end
+
+    batch_size = 4
+    data_length = 16
+    item_size = (28, 28, 1)
+
+    dataset = ZerosDataset{Float32, 4}(item_size, data_length)
+    loader = DataLoader(dataset, batchsize=batch_size, shuffle=true)
+
+    @inferred first(loader)
+    @test length(loader) == data_length / batch_size
+
+    batches = collect(loader)
+    @test length(batches) == data_length / batch_size
+
+    for b in batches
+        @test size(b) == (item_size..., batch_size)
+    end
+end
+
 @testset "CMUDict" begin 
     @test cmudict()["CATASTROPHE"] == :[K,AH0,T,AE1,S,T,R,AH0,F,IY0].args