# Broadcasting over a Flux.DataLoader

**URL:** https://discourse.julialang.org/t/broadcasting-over-a-flux-dataloader/70414
**Category:** Machine Learning
**Created:** [October 26, 2021, 4:20pm UTC](https://discourse.julialang.org/t/broadcasting-over-a-flux-dataloader/70414 "2021-10-26T16:20:30Z")
**Posts on this page:** 12
**Page:** 1

<div class="post-metadata">

### Author: ![austinbean](https://avatars.discourse-cdn.com/v4/letter/a/839c29/32.png) [@austinbean](https://discourse.julialang.org/u/austinbean)
#### Post date: [October 26, 2021, 4:20pm UTC](https://discourse.julialang.org/t/broadcasting-over-a-flux-dataloader/70414/1 "2021-10-26T16:20:30Z")

</div>

Apologies for this question since I’ve had a lot of trouble producing a MWE (which I don’t have).

I’m working on an RNN using text data (each observation is a sentence - each word in each sentence is mapped to a word embedding). The problem is classification into six classes. The issue I am having I think relates to using the Flux.DataLoader type.

The basic setup is…

```julia
train_data = Flux.DataLoader((data=train, label=tlabels), batchsize=42, shuffle=true)

# train_data.data.data is 1x357, where each item in the vector is (vocab_size)x(sentence_length) and is a matrix of one-hot vectors
# train_data.data.label is 1x357, where each item in the vector is (6,) and is one-hot encoded

function two_layers(args)
    # Embed is a layer I defined to take one-hot vectors to word embeddings of length N=50
    scanner = Chain(Embed(args.vocab, args.embed_len, args.emb_table), LSTM(args.embed_len, args.N), LSTM(args.N, args.N))
    encoder = Dense(args.N, args.classes, identity)
	return scanner, encoder 
end

function model(x, scanner, encoder)
	state = scanner(x)[:,end]     
	reset!(scanner)                   
	encoder(state)                      
end 

ps = params(scanner, encoder)   

loss(x,y) = Flux.logitcrossentropy(model(x, scanner, encoder), y)

```

When I do:

```julia
	for (x,y) in train_data 
		gs = Flux.gradient(()->loss(x,y), ps)
		Flux.update!(opt, ps, gs)
	end 

```

I get the error:

```julia
ERROR: LoadError: MethodError: no method matching (::Flux.LSTMCell{Matrix{Float32}, Vector{Float32}, Tuple{Matrix{Float32}, Matrix{Float32}}})(::Tuple{Matrix{Float32}, Matrix{Float32}}, ::Matrix{Real})

```

I this error find confusing b/c I don’t know what `::Matrix{Real}` is being referred to at the end of the error message.

I expected `(x,y)` to iterate over batches of data and labels (which printing some type information about `x, y` inside the loop suggests it is doing).

```julia
for (x,y) in train_data 
		@info "type of (x,y)" typeof((x,y))
		@info "type of x", typeof(x)
		@info "size x", size(x)
		@info "size x[1]", size(x[1])
		@info "typeof y", typeof(y)
		@info "size y", size(y)
		@info "size y[1]", size(y[1])
end

Info: type of (x,y)
└ typeof((x, y)) = Tuple{Matrix{Flux.OneHotArray{UInt32, 733, 1, 2, Vector{UInt32}}}, Matrix{Flux.OneHotArray{UInt32, 6, 0, 1, UInt32}}}
[ Info: ("type of x", Matrix{Flux.OneHotArray{UInt32, 733, 1, 2, Vector{UInt32}}})
[ Info: ("size x", (1, 42)) # 42 is the batch size
[Info: ("size x[1]", (733, 18)) # 18 is the length of the sentence
[ Info: ("typeof y", Matrix{Flux.OneHotArray{UInt32, 6, 0, 1, UInt32}})
[ Info: ("size y", (1, 42))
[Info: ("size y[1]", (6,))

```

I also can do `loss(x[1], y[1])` and get a number. So I think my problems might be solved by…

1. Figuring out what that extra `::Matrix{Real}` is and where it comes from?
2. Broadcasting over the DataLoader type? (I have tried ` gradient(()->loss.(x,y)` but that also gives an error.)

I’m definitely missing something but I can’t figure out what! I appreciate any help and apologize for not having an MWE…

---

<div class="post-metadata">

### Author: ![ToucheSir](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/touchesir/32/14411_2.png) [@ToucheSir](https://discourse.julialang.org/u/ToucheSir)
#### Post date: [October 28, 2021, 4:07pm UTC](https://discourse.julialang.org/t/broadcasting-over-a-flux-dataloader/70414/2 "2021-10-28T16:07:48Z")

</div>

I’d recommend having a read through [Model Reference · Flux](https://fluxml.ai/Flux.jl/stable/models/layers/#Recurrent-Layers) because it describes what data formats are supported by Flux’s RNNs. In short, a DataLoader can give you a batch of sequences, but that must be converted to a sequence of batches in order to work with RNNs.

---

<div class="post-metadata">

### Author: ![austinbean](https://avatars.discourse-cdn.com/v4/letter/a/839c29/32.png) [@austinbean](https://discourse.julialang.org/u/austinbean)
#### Post date: [October 28, 2021, 7:06pm UTC](https://discourse.julialang.org/t/broadcasting-over-a-flux-dataloader/70414/3 "2021-10-28T19:06:39Z")

</div>

Thanks very much for your reply.

I am still very puzzled after reading the docs again. I’m working on a simplified example but I cannot make it work. Maybe you can help me get the simple case working and then I can figure out my actual use case?

```julia
# Suppose I have four sentences, each 20 words long, 
# w/ embedding dim per word of 10

a1 = rand(Float32, 10,20)
a2 = rand(Float32, 10,20)
a3 = rand(Float32, 10,20)
a4 = rand(Float32, 10,20)
fake_train = cat(a1, a2, a3, a4; dims = 3)

fake_labels = Flux.onehotbatch([0,1,0,1], [0,1])
# First question: is this step necessary to give the labels a similar shape to fake_train??
fake_labels = reshape(fake_labels, (2,1,4)) 

DL = Flux.DataLoader((data = fake_train, label=fake_labels), batchsize=2)

l1 = LSTM(10, 5)
d1 = Dense(5,2)
function model(x, scanner, encoder)
    state = scanner(x)[:,end]
    reset!(scanner)
    encoder(state)
end 

ps = params(l1, d1)
opt = ADAM(1e-3)

loss(x,y)= Flux.logitbinarycrossentropy(model(x, l1, d1), y)
# returns a number, not an error, so the loss function is working, I believe
loss(fake_train[:,:,1], fake_labels[:,:,1]) 

for i = 1:3
    @info i 
    Flux.train!(loss, ps, DL, opt)
end 

ERROR: LoadError: MethodError: no method matching loss(::NamedTuple{(:data, :label), Tuple{Array{Float32, 3}, Array{Bool, 3}}})
# this error makes sense given the way I defined loss(x,y)

```

The docs [say](https://fluxml.ai/Flux.jl/stable/training/training/#Flux.Optimise.train!) “If `d` is a tuple of arguments to `loss` call `loss(d...)` , else call `loss(d)` .” I tried to follow advice [here](https://discourse.julialang.org/t/flux-jl-methoderror-no-method-matching-loss/63478/3) and do literally that (i.e., `train!(x -> loss(x...), ...)`, but that also generated an error:  
`ERROR: LoadError: BoundsError: attempt to access 5×20×2 Array{Float32, 3} at index [1:5, 20]`

I have also tried defining `loss(x,y) = logitbinarycrossentropy.(...` (i.e, broadcasting) or `loss((x,y)) = ...` (i.e., defining the loss over tuples directly), but these generate errors of their own.

Maybe I am a little closer to an MWE now…?

---

<div class="post-metadata">

### Author: ![ToucheSir](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/touchesir/32/14411_2.png) [@ToucheSir](https://discourse.julialang.org/u/ToucheSir)
#### Post date: [October 28, 2021, 8:32pm UTC](https://discourse.julialang.org/t/broadcasting-over-a-flux-dataloader/70414/4 "2021-10-28T20:32:56Z")

</div>

Ah sorry, apparently I posted the wrong docs page! Strongly recommend reading through [Recurrence · Flux](https://fluxml.ai/Flux.jl/stable/models/recurrence/). It’ll tell you pretty much all you need to know to get an RNN working. Then the remaining part is getting your data into a form which is RNN-friendly. Here are some tips for that:

1. Get one batch working first. e.g. write functionality to take a batch of sequences and transform them into a (possibly padded) sequence of batched features. This includes making sure the loss function works as expected.
2. Once a single batch works, you can try using a dataloader or other functionality to sample batches from the training set.
3. `train!` is both very high level and not very flexible. As a consequence, it’s basically a no-go for RNNs. Since you’re comfortable with a custom training loop already, that’s the way to go.

---

<div class="post-metadata">

### Author: ![austinbean](https://avatars.discourse-cdn.com/v4/letter/a/839c29/32.png) [@austinbean](https://discourse.julialang.org/u/austinbean)
#### Post date: [October 29, 2021, 2:12am UTC](https://discourse.julialang.org/t/broadcasting-over-a-flux-dataloader/70414/5 "2021-10-29T02:12:55Z")

</div>

Thanks for commenting again.

I can make the loss function over a batch one element at a time, but not with the `DataLoader` type.

Example:

```julia
using Flux
using Flux: reset!, onehotbatch 

fake_train = rand(Float32,10,20, 4)
fake_labels = Flux.onehotbatch([0,1,0,1], [0,1])
 
l1 = LSTM(10, 5)
d1 = Dense(5,2)
function model(x, scanner, encoder)
    state = scanner(x)[:,end]
    reset!(scanner)
    encoder(state)
end 

ps = params(l1, d1)
opt = ADAM(1e-3)

loss(x,y) = Flux.logitbinarycrossentropy(model(x, l1, d1), y) 

# this works
for (x,y) in zip(eachslice(fake_train, dims=3), fake_labels)
    grads = Flux.gradient(ps) do 
        loss(x,y)
    end
    Flux.Optimise.update!(opt, ps, grads)
end
 

```

But if I try to extend the loss function to handle a batch via the `DataLoader` (more or less [following the docs you linked to](https://fluxml.ai/Flux.jl/stable/models/recurrence/)) I get an error I can’t figure out:

```julia
DL = Flux.DataLoader((data = fake_train, label=fake_labels), batchsize=2)

function loss2(x,y)
    sum( Flux.logitbinarycrossentropy(model(dat, l1, d1), lab ) for (dat,lab) in zip(x,y))
end 

for (x,y) in DL
    loss2(x,y) 
end 

ERROR: LoadError: MethodError: no method matching (::Flux.LSTMCell{Matrix{Float32}, Vector{Float32}, Tuple{Matrix{Float32}, Matrix{Float32}}})(::Tuple{Matrix{Float32}, Matrix{Float32}}, ::Float32)

```

The part I don’t understand about that error is the single `::Float32` at the end of `(::Tuple{Matrix{Float32}, Matrix{Float32}}, ::Float32)`. What `Float32` is that?

I expect that `x` should be a `Matrix{Float32}` with size `(10,20,2)` and that `y` would be a onehot matrix with size `(2,2,2)`. If `DataLoader` treats the last dimension as the observation dimension ([right](https://fluxml.ai/Flux.jl/stable/data/dataloader/)?), it should split `fake_labels` like that.

Maybe `fake_labels` is getting mangled somehow? I do not know if that is the problem though.

---

<div class="post-metadata">

### Author: ![ToucheSir](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/touchesir/32/14411_2.png) [@ToucheSir](https://discourse.julialang.org/u/ToucheSir)
#### Post date: [November 3, 2021, 2:08am UTC](https://discourse.julialang.org/t/broadcasting-over-a-flux-dataloader/70414/6 "2021-11-03T02:08:49Z")

</div>

Let’s work backwards to see what is going on, starting from the forward pass:

```julia
function model(x, scanner, encoder)
    state = scanner(x)[:,end]
    reset!(scanner)
    encoder(state)
end 

```

1. Per the linked docs, what structure should `x` have? If it’s not a sequence of `num features x batch size` arrays, then the input format is incorrect.
2. Given an input of the correct shape, how should `scanner` be applied over them? Here you’ll want to consult [Recurrence · Flux](https://fluxml.ai/Flux.jl/stable/models/recurrence/#Working-with-sequences) to see what the correct usage pattern is.
3. Once the RNN part is correct, how should the RNN output be fed into the dense encoder? Your code is already pretty close there, but I posit to you that this should be even easier once 2) is figured out.

That’s the entirety of the forward pass. Once that’s working as expected, you can move onto making sure the loss function inputs and outputs are what you expect. I recommend starting with a dummy batch of data first before trying to integrate data loading, as it’ll remove `Flux.DataLoader` as a variable.

---

<div class="post-metadata">

### Author: ![austinbean](https://avatars.discourse-cdn.com/v4/letter/a/839c29/32.png) [@austinbean](https://discourse.julialang.org/u/austinbean)
#### Post date: [November 3, 2021, 8:47pm UTC](https://discourse.julialang.org/t/broadcasting-over-a-flux-dataloader/70414/7 "2021-11-03T20:47:44Z")

</div>

I thought I had it, but now I’m pretty sure I don’t! I appreciate your taking the time to look at this again.

[This part](https://fluxml.ai/Flux.jl/stable/models/recurrence/) (which you are highlighting again in your most recent comment) I think I’ve got:

> A potential source of ambiguity with RNN in Flux can come from the different data layout compared to some common frameworks where data is typically a 3 dimensional array: `(features, seq length, samples)` . In Flux, those 3 dimensions are provided through a vector of seq length containing a matrix `(features, samples)` .

I have sentences of max length 80, an embedding dimension of 50, for (let’s say) 500 training samples. Then `training_data` is a `Vector{ Matrix{Float32}}` where `size(training_data) = (80,)` and `size(data[1]) = (50,500)`. `train_label` is `(1,500)` where each entry is a one-hot encoded label for `k` classes in the problem (I think this size actually won’t work later).

I’m sure there was a deliberate, careful decision behind this data layout, but it did not correspond well to the way my data was organized (as `(sentence, label)` rows in a spreadsheet).

Now suppose I do:

```julia
l1 = LSTM(50, 5) # 50 is the embedding dimension

```

Given the data organization above, the _last_ feature for each observation is actually given by:

```julia
l1.(train_data)[end] # not [:, end]

```

So what I want is roughly:

```julia
d = Dense(5,k)
d(l1.(train_data)[end]) # broadcast l1, but not d
# gives correctly sized output k x (# sentences)

```

I think this gets it:

```julia
function model(x, scanner, encoder)
	state = scanner.(x)[end] # the last row
	reset!(scanner) # will this be called N_sentences times?
	encoder(state)
end 

```

I still need to understand how to make the `loss` work. My data is not shaped correctly. I will make another post when I can figure that out… Again I appreciate your help.

Two comments only because you are a committer to the repository!

It’s hard to test to test the `model(x)` functions with data organized as a `Vector{ Matrix{Float32}}` because it’s hard to select out “one observation”. I can’t do `model(data[:][:,1])`, for example, where `data[:][:,1]` would be all the embeddings for all of the words in the first sentence.

Another problem: isn’t it a little hard to batch training examples with this data layout? AFAICT, to batch this, you have to produce `Vectors` of length `max_length` (=80) where each element is a matrix w/ `features` rows (=50) and with `batch_size` columns. This is easier if the data is stored as a 3-D tensor but it is a little harder when it’s a `Vector{ Matrix{Float32} }`. Or is there an easier way?

---

<div class="post-metadata">

### Author: ![ToucheSir](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/touchesir/32/14411_2.png) [@ToucheSir](https://discourse.julialang.org/u/ToucheSir)
#### Post date: [November 3, 2021, 11:18pm UTC](https://discourse.julialang.org/t/broadcasting-over-a-flux-dataloader/70414/8 "2021-11-03T23:18:09Z")

</div>

> [@austinbean](#):
>
> I have sentences of max length 80, an embedding dimension of 50, for (let’s say) 500 training samples. Then `training_data` is a `Vector{ Matrix{Float32}}` where `size(training_data) = (80,)` and `size(data[1]) = (50,500)` . `train_label` is `(1,500)` where each entry is a one-hot encoded label for `k` classes in the problem (I think this size actually won’t work later).
> 
> I’m sure there was a deliberate, careful decision behind this data layout, but it did not correspond well to the way my data was organized (as `(sentence, label)` rows in a spreadsheet).

Remember that input dimensions in Julia are generally reversed from many Python libraries. So if in PyTorch you’d pass something of size `(80, 500, 50)` to a [RNN — PyTorch 2.1 documentation](https://pytorch.org/docs/stable/generated/torch.nn.RNN.html), in Flux that’d be `(50, 500, 80)` (or more accurately, `((50, 500), 80)`, more on that in a sec). The easiest way to think about this is that the sequence length gets added as an adjacent dimension to the batch size, i.e. `batch x features` becomes `length x batch x features` and `features x batch` becomes `features x batch x length`.

> [@austinbean](#):
>
> Given the data organization above, the _last_ feature for each observation is actually given by:
> 
> ```julia
> l1.(train_data)[end] # not [:, end]
> 
> ```

Yes, but note the warning under [Recurrence · Flux](https://fluxml.ai/Flux.jl/stable/models/recurrence/#Working-with-sequences) :

> Mapping and broadcasting operations with stateful layers such are discouraged, since the julia language doesn’t guarantee a specific execution order…

Since you’re only interested in the final output, there’s also no need to materialize any of the intermediate timesteps either:

```julia
function model(x, scanner, encoder)
    # this will be called once, but calling it in the middle of the forward pass may not work. Best to call it outside of the loss function completely, 
    # but if that's not possible then moving it here is better.
    reset!(scanner) 

    state = x[1]
    for i in 2:length(x)
        state = scanner(x[i])
    end
    encoder(state)
end 

```

`foldl` may be substituted for a more functional forwards pass.

---

<div class="post-metadata">

### Author: ![ToucheSir](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/touchesir/32/14411_2.png) [@ToucheSir](https://discourse.julialang.org/u/ToucheSir)
#### Post date: [November 3, 2021, 11:49pm UTC](https://discourse.julialang.org/t/broadcasting-over-a-flux-dataloader/70414/9 "2021-11-03T23:49:50Z")

</div>

> [@austinbean](#):
>
> It’s hard to test to test the `model(x)` functions with data organized as a `Vector{ Matrix{Float32}}` because it’s hard to select out “one observation”. I can’t do `model(data[:][:,1])` , for example, where `data[:][:,1]` would be all the embeddings for all of the words in the first sentence.

This is a sore point of the current RNN interface, and to your earlier point I’m not really sure why it was designed this way. We’ve been slowly moving towards supporting contiguous sequences of `AbstractArray{T, 3}` (see [https://github.com/FluxML/Flux.jl/pull/1686](https://github.com/FluxML/Flux.jl/pull/1686), now in stable Flux. Also the tracking issue at [Recurrent network interface updates/design · Issue #1678 · FluxML/Flux.jl · GitHub](https://github.com/FluxML/Flux.jl/issues/1678)) like other frameworks. So your example would become `rnn_model(data[:, 1, :])` (or `rnn_model(@view data[:, 1, :]` to avoid the extra allocation).

> [@austinbean](#):
>
> Another problem: isn’t it a little hard to batch training examples with this data layout? AFAICT, to batch this, you have to produce `Vectors` of length `max_length` (=80) where each element is a matrix w/ `features` rows (=50) and with `batch_size` columns. This is easier if the data is stored as a 3-D tensor but it is a little harder when it’s a `Vector{ Matrix{Float32} }` . Or is there an easier way?

If you can create a 3D array, then converting it to a `Vector{Matrix}` is a simple matter of `collect(eachslice(x, dims=3))` 🙂

---

<div class="post-metadata">

### Author: ![austinbean](https://avatars.discourse-cdn.com/v4/letter/a/839c29/32.png) [@austinbean](https://discourse.julialang.org/u/austinbean)
#### Post date: [November 4, 2021, 3:24pm UTC](https://discourse.julialang.org/t/broadcasting-over-a-flux-dataloader/70414/10 "2021-11-04T15:24:05Z")

</div>

Good to know this may change! Thanks for continuing to look at this question.

So having changed the model to this:

```julia
function model(x, scanner, encoder)
	reset!(scanner)                  
	state = scanner(x[1]) 
	for i=2:length(x) # this is explicit about the order
		state = scanner(x[i])
	end 
	encoder(state)
end

```

The only way I could make the loss function work across a batch of `data` organized as above and `labels` was like this:

```julia
function overall_loss(data, labels, scanner, encoder)
	losses = zeros(Float32, length(labels))
	for i = 1:size(losses,1)
		losses[i] += Flux.logitcrossentropy(model([data[j][:,i] for j = 1:size(data,1)] , scanner, encoder), labels[i])
	end 
	return mean(losses)
end 

```

That’s very slow (about 2:30 for one evaluation on a batch of 350 sentences), (I think) due to the time it takes to collect `[data[j][:,i] for j = 1:size(data,1)]`.

Slightly faster was to collect all of those vectors in advance:

```julia
function another_loss(data,labels,scanner,encoder)
	losses = zeros(Float32,length(labels))
	td = [[data[i][:,k] for i =1:length(data)] for k = 1:size(data[1],2)]
	for i = 1:size(td,1)
		losses[i] += Flux.logitcrossentropy(model(td[i], scanner, encoder) ,labels[i])
	end 
	return mean(losses)
end 

```

But both are very slow. Is there a smarter way to do that?

---

<div class="post-metadata">

### Author: ![ToucheSir](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/touchesir/32/14411_2.png) [@ToucheSir](https://discourse.julialang.org/u/ToucheSir)
#### Post date: [November 4, 2021, 5:26pm UTC](https://discourse.julialang.org/t/broadcasting-over-a-flux-dataloader/70414/11 "2021-11-04T17:26:41Z")

</div>

`logitcrossentropy` and all other Flux losses should already be vectorized, so if the following doesn’t work then I’d check the shape of the labels and model outputs to make sure they’re the correct shape (`output features x batch size`):

```julia
function overall_loss(data, labels, scanner, encoder)
  logits = model(data, scanner, encoder)
  Flux.Zygote.ignore() do # Debugging purposes only. If this fails, check the shapes!
    @assert ndims(logits) == 2 && size(logits) == size(labels)
  end
  Flux.logitcrossentropy(logits, labels)
end

```

---

<div class="post-metadata">

### Author: ![austinbean](https://avatars.discourse-cdn.com/v4/letter/a/839c29/32.png) [@austinbean](https://discourse.julialang.org/u/austinbean)
#### Post date: [November 4, 2021, 6:59pm UTC](https://discourse.julialang.org/t/broadcasting-over-a-flux-dataloader/70414/12 "2021-11-04T18:59:00Z")

</div>

I swear `model(data, scanner, encoder)` was giving me an error yesterday evening, but it works as you propose here (after reshaping my labels).

Indeed your version is much faster: 35 seconds.

Maybe the entire thing actually works now…
