Flux: ERROR: OutOfMemoryError()


I am new to Julia (3-4 days experience only) and trying to fit a simple least square neural network model in Julia 1.2.0 using Flux. The code at the end results in ERROR: OutOfMemoryError() even though I have 64 GB of memory. However, if I let n = 10000 it works, but I do not understand the memory usage

julia> @time Flux.train!(loss, params(model), data_iterator, ADAM())
 13.316979 seconds (15.39 M allocations: 26.925 GiB, 9.88% gc time)

Any hints?

using Flux

n = 100000
p = 10
x = rand(p, n)  
y = rand(n)

model = Chain(Dense(p, 32, relu),
              Dense(32, 16, relu),
              Dense(16, 1))
loss(x, y) = Flux.mse(model(x), y)
data_iterator = Iterators.repeated((x, y), 5)
@time Flux.train!(loss, params(model), data_iterator, ADAM())

I have realized that my example is not a very realistic one. Normally, one would use a batch size of, say, 100 and in that case it works well.

trdata = [(rand(p,100), rand(100)) for i in 1:1000 ]
@time Flux.@epochs 5 Flux.train!(loss, params(model), trdata, ADAM())