search

JuliaML / MLUtils.jl

Utilities and abstractions for Machine Learning tasks
MIT License
107 stars 20 forks source link

Padded batching #78

Open casper2002casper opened 2 years ago

casper2002casper commented 2 years ago

When working with ANN that output different sized vectors depending on the input (for example when using GraphNeuralNetworks.jl), it would be useful to convert the output of a batch to a CuArray in order to perform loss computations. Current:

julia> MLUtils.batch([[1,2],[3,4]])
2×2 Matrix{Int64}:
 1  3
 2  4

julia> MLUtils.batch([[1,2],[3]])
ERROR: DimensionMismatch("mismatch in dimension 1 (expected 2 got 1)")

Feature:

julia> MLUtils.batch([[1,2],[3]], pad =  0)
2×2 Matrix{Int64}:
 1  3
 2  0
darsnack commented 2 years ago

As a temporary workaround, you can do

batches = [[1, 2], [3]]
MLUtils.batch(rpad.(batches, 2, 0))

Probably a variation of this where we pad as we iterate batches would be a possible PR for this feature.