refactor
This commit is contained in:
parent
1bd0a43b7d
commit
0ed4e65d2f
@ -15,7 +15,13 @@ size(b::CatMat) = (size(b.data, 1),)
|
||||
|
||||
getindex(b::CatMat, i)::eltype(b) = slicedim(b.data, 1, i)
|
||||
|
||||
setindex!(b::CatMat, v, i) = b.data[i, :] = v
|
||||
setindex!(b::CatMat, v, i::Integer) = b.data[i, :] = v
|
||||
|
||||
function setindex!(b::CatMat, xs, ::Colon)
|
||||
for (i, x) in enumerate(xs)
|
||||
b[i] = x
|
||||
end
|
||||
end
|
||||
|
||||
allequal(xs) = all(x -> x == first(xs), xs)
|
||||
|
||||
|
114
src/dims/iter.jl
114
src/dims/iter.jl
@ -1,72 +1,72 @@
|
||||
export Batched
|
||||
|
||||
zipt(xs...) = (xs,)
|
||||
zipt(xs::Tuple...) = zip(xs...)
|
||||
|
||||
import Base: start, next, done, iteratorsize, iteratoreltype, eltype, length
|
||||
|
||||
mutable struct Batched{T,S}
|
||||
batch::Int
|
||||
iter::T
|
||||
"`Batched` always read a batch in advance, and store it in `buf`"
|
||||
# Stateful iteration
|
||||
|
||||
mutable struct StatefulIter{I,S,T}
|
||||
iter::I
|
||||
state::S
|
||||
next::Nullable{T}
|
||||
end
|
||||
|
||||
function StatefulIter(itr)
|
||||
state = start(itr)
|
||||
val, state = done(itr, state) ? (Nullable(), state) : next(itr, state)
|
||||
return StatefulIter(itr, state, convert(Nullable, val))
|
||||
end
|
||||
|
||||
peek(s::StatefulIter) = get(s.next)
|
||||
|
||||
function Base.take!(s::StatefulIter)
|
||||
x = peek(s)
|
||||
if !done(s.iter, s.state)
|
||||
s.next, s.state = next(s.iter, s.state)
|
||||
else
|
||||
s.next = Nullable()
|
||||
end
|
||||
return x
|
||||
end
|
||||
|
||||
Base.isempty(s::StatefulIter) = isnull(s.next)
|
||||
Base.eltype(s::StatefulIter) = eltype(s.next)
|
||||
|
||||
function taken!(s::StatefulIter, n::Integer)
|
||||
xs = eltype(s)[]
|
||||
for _ = 1:n
|
||||
isempty(s) && break
|
||||
push!(xs, take!(s))
|
||||
end
|
||||
return xs
|
||||
end
|
||||
|
||||
# Batched
|
||||
|
||||
struct Batched{I<:StatefulIter,S}
|
||||
itr::I
|
||||
buf::S
|
||||
i
|
||||
end
|
||||
|
||||
function Batched(iter::T, batch::Integer) where T
|
||||
batch >= 1 || throw(ArgumentError("batch size must >= 1"))
|
||||
i = start(iter)
|
||||
done(iter, i) && return Batched{T,Void}(batch, iter, nothing, i)
|
||||
v, i = next(iter, i)
|
||||
|
||||
buf = mapt(v) do x
|
||||
storage = Array{eltype(x)}(batch, size(x)...)
|
||||
storage[1, :] = x
|
||||
rebatch(storage)
|
||||
end
|
||||
|
||||
for ibatch in 2:batch
|
||||
if done(iter, i)
|
||||
warn("data less than one batch will be ignored, please use a smaller batch size")
|
||||
return Batched{T,Void}(batch, iter, nothing, i)
|
||||
end
|
||||
|
||||
v, i = next(iter, i)
|
||||
map(x->setindex!(x..., ibatch), zipt(buf, v))
|
||||
end
|
||||
|
||||
Batched{T,typeof(buf)}(batch, iter, buf, i)
|
||||
function Batched(itr, n::Integer)
|
||||
n >= 1 || throw(ArgumentError("batch size must be >= 1"))
|
||||
itr = StatefulIter(itr)
|
||||
buf = rebatch(similar(eltype(itr)(), n, size(peek(itr))...))
|
||||
Batched(itr, buf)
|
||||
end
|
||||
|
||||
iteratoreltype(::Type{Batched{T,S}}) where {T,S} = Base.HasEltype()
|
||||
iteratoreltype(::Type{<:Batched}) = Base.HasEltype()
|
||||
iteratorsize(::Type{<:Batched}) = Base.SizeUnknown()
|
||||
|
||||
iteratorsize(::Type{Batched{T,S}}) where {T,S} =
|
||||
iteratorsize(T) isa Base.HasShape ?
|
||||
Base.HasLength() : iteratorsize(T)
|
||||
eltype{T,S}(x::Batched{T,S}) = S
|
||||
|
||||
length(x::Batched) = length(x.iter) ÷ x.batch
|
||||
start(::Batched) = ()
|
||||
|
||||
eltype(x::Batched{T,S}) where {T,S} = S
|
||||
|
||||
start(x::Batched) = true
|
||||
|
||||
next(x::Batched, ::Bool) = x.buf, false
|
||||
next(x::Batched, _) = x.buf, ()
|
||||
|
||||
# will be less hacky if https://github.com/JuliaLang/julia/issues/18823
|
||||
function done(x::Batched, fresh)
|
||||
fresh && return false
|
||||
|
||||
for ibatch in 1:x.batch
|
||||
if done(x.iter, x.i)
|
||||
ibatch != 1 && warn("cannot perfectly divide data by batch size, remainder will be discarded")
|
||||
return true
|
||||
end
|
||||
|
||||
v, x.i = next(x.iter, x.i)
|
||||
map(x->setindex!(x..., ibatch), zipt(x.buf, v))
|
||||
end
|
||||
|
||||
false
|
||||
function done(x::Batched, _)
|
||||
next = taken!(x.itr, length(x.buf))
|
||||
length(next) < length(x.buf) && return true
|
||||
x.buf[:] = next
|
||||
return false
|
||||
end
|
||||
|
||||
done(::Batched{T,Void}, ::Bool) where T = true
|
||||
|
Loading…
Reference in New Issue
Block a user