Flux.jl: How to train a batch of data

[pylat]

I have a basic question but I couldn't understand from the docs how exactly you can use a batch of data rather than a single data point. To clarify, consider this simple example with gradient descent with stepsize η=0.1 for a two class classification:

xs = [x_class1; x_class2];  # two classes of data
ys = [fill(0,size(x_class1)); fill(1,size(x_class2))];  # labels
#######      gradient descent step on one data only 
model = Dense(2,1,Flux.relu)
η = 0.1
for i in 1:100
loss = Flux.mse(model(xs[1]),ys[1])   ## training with one sample
back!(loss)
model.W.data .-= model.W.grad * η
model.b.data .-= model.b.grad * η
end

How would you modify this code to take the full gradient step for several data points rather than one (the gradient being sum of each individual gradient)?
The same question applies to the Flux.Train!: How should the input data be for a batch update?

[MikeInnes]

You just need to stack a bunch of inputs into a single array (batch dimension last) which most layers will handle. For example rand(5) represents an input with 5 features; rand(5, 3) represents 3 such inputs. If you pass that 2D array into a dense layer like m(x), you'll see that each column of it gets processed independently, as if you had written m(x[:, 1]) etc. Loss functions are also aware of 2D arrays, so things like mse will give you the sum of losses across the batch; then the gradient you get will be the sum of all the independent gradients.

[pylat]

Nice. Thanks.

[jklaise]

Was just trying to figure this out myself, I feel this (i.e. batch dimension last convention) should be prominent in the getting started part of the documentation.