# Flux.jl vanilla ANN loss goes to NaN with mini batch

**URL:** <https://discourse.julialang.org/t/flux-jl-vanilla-ann-loss-goes-to-nan-with-mini-batch/25511>\
**Category:** Machine Learning\
**Tags:** question\
**Created:** [June 21, 2019, 4:14am UTC](https://discourse.julialang.org/t/flux-jl-vanilla-ann-loss-goes-to-nan-with-mini-batch/25511 "2019-06-21T04:14:00Z")\
**Posts on this page:** 1\
**Page:** 1

<div class="post-metadata">

**Author:** ![jchen975](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/jchen975/32/8994_2.png) [@jchen975](https://discourse.julialang.org/u/jchen975)\
**Post date:** [June 21, 2019, 4:14am UTC](https://discourse.julialang.org/t/flux-jl-vanilla-ann-loss-goes-to-nan-with-mini-batch/25511/1 "2019-06-21T04:14:00Z")

</div>

I’m building a vanilla NN with 1 or 2 hidden layers for a regression problem that’s basically as follows:  
I take N sets of known values (x\_1, y\_1), ..., (x\_k, y\_k) as input to an approximation solver that returns N sets of (a\_1, b\_1), ..., (a\_k, b\_k), and to a solver for true values that returns N sets of (c\_1, d\_1), ..., (c\_k, d\_k), where a\_i \approx c\_i and b\_i \approx d\_i, for i = 1, 2, ..., k, but could have large inaccuracies. I’m building an ANN that takes N x (k\*4) input [x\_1, y\_1, a\_1, b\_1, x\_2, y\_2, a\_2, b\_2,...] and outputs [c\_1', d\_1', c\_2', d\_2', ...] that are better approximations to the true c, d values than the original a, b approximations.

I’m using ADAM optimizer, and naturally, my loss function is `MSELoss` and the model is simply two `Dense` layers chained together. All worked well, except when I have a large input size where K\*4 \approx 6000 and N = 10000, training the full batch on GPU (Nvidia GTX960M, pretty trash but that’s all I have for now) gave me `OutOfMemoryError()`. The maximum N I can run full batch training without the error is about 2500. So I decided to try mini-batching the input, and feed the mini batches to GPU in each iteration as following:

```julia
randIdx = collect(1:1:size(trainData)[2]) # trainData has shape K*4 x N
numBatches = round(Int, floor(size(trainData)[2] / batch_size))
for epoch = 1:epochs
	println("epoch: ", epoch)
	Random.shuffle!(randIdx) # to shuffle training set
	i = 1  
	for j = 1:numBatches
		println(j)
		batchData = trainData[:, randIdx[i:batch_size]] |> gpu
		batchTarget = trainTarget[:, randIdx[i:batch_size]] |> gpu
		Flux.train!(loss, Flux.params(model), [(batchData, batchTarget)], opt)
		epochTrainLoss += Tracker.data(loss(batchData, batchTarget))
		epochTrainAcc += Tracker.data(accuracy(batchData, batchTarget))
		i += batch_size
	end
	push!(trainLoss, epochTrainLoss / numBatches)
	push!(trainAcc, epochTrainAcc / numBatches)
	push!(valLoss, Tracker.data(loss(valData, valTarget)))
	push!(valAcc, Tracker.data(accuracy(valData, valTarget)))
	epochTrainLoss, epochTrainAcc = 0.0, 0.0 # reset values
end

```

However, now in the second iteration of the first epoch, I’m getting `Loss is NaN` error:

```julia
julia> include("case_general.jl")
epoch: 1
1
2.11837100982666
2
ERROR: LoadError: Loss is NaN
Stacktrace:
 [1] losscheck(::Tracker.TrackedReal{Float32}) at C:\Users\me\.julia\packages\Tracker\RRYy6\src\back.jl:155
 [2] gradient_(::getfield(Flux.Optimise, Symbol("##14#20")){getfield(Main, Symbol("#loss#256"))}, ::Tracker.Params) at C:\Users\me\.julia\packages\Tracker\RRYy6\src\back.jl:98
 [3] #gradient#24(::Bool, ::Function, ::Function, ::Tracker.Params) at C:\Users\me\.julia\packages\Tracker\RRYy6\src\back.jl:164
 [4] gradient at C:\Users\me\.julia\packages\Tracker\RRYy6\src\back.jl:164 [inlined]
 [5] macro expansion at C:\Users\me\.julia\packages\Flux\qXNjB\src\optimise\train.jl:71 [inlined]
 [6] macro expansion at C:\Users\me\.julia\packages\Juno\TfNYn\src\progress.jl:124 [inlined]
 [7] #train!#12(::getfield(Flux.Optimise, Symbol("##16#22")), ::Function, ::Function, ::Tracker.Params, ::Array{Tuple{CuArray{Float32,2},CuArray{Float32,2}},1}, ::ADAM) at C:\Users\me\.julia\package\train.jl:69
 [8] train!(::Function, ::Tracker.Params, ::Array{Tuple{CuArray{Float32,2},CuArray{Float32,2}},1}, ::ADAM) at C:\Users\me\.julia\packages\Flux\qXNjB\src\optimise\train.jl:67
 [9] #mlp#254(::Int64, ::Function, ::String, ::Array{Float64,2}, ::Array{Float64,2}, ::Float64, ::Int64, ::Int64, ::Int64) at F:\work\large cases\case_general.jl:156
 [10] (::getfield(Main, Symbol("#kw##mlp")))(::NamedTuple{(:K2,),Tuple{Int64}}, ::typeof(mlp), ::String, ::Array{Float64,2}, ::Array{Float64,2}, ::Float64, ::Int64, ::Int64, ::Int64) at .\none:0
 [11] macro expansion at .\util.jl:156 [inlined]
 [12] main(::Int64) at F:\work\large cases\case_general.jl:221
 [13] top-level scope at none:0
 [14] include at .\boot.jl:326 [inlined]
 [15] include_relative(::Module, ::String) at .\loading.jl:1038
 [16] include(::Module, ::String) at .\sysimg.jl:29
 [17] include(::String) at .\client.jl:403
 [18] top-level scope at none:0
in expression starting at F:\work\large cases\case_general.jl:232

```

I looked at [this link on NaN loss](https://discourse.julialang.org/t/how-come-flux-jls-network-parameters-go-to-nan/16439) and tried adjusting my learning rate to extremely small (1e-15) and increasing my batch size to 2500 (so that it should have the same behavior as full batch training with 2500 samples), but those didn’t help, and only the first iteration loss slightly changed (around 2.1). I also checked my input data and target arrays and made sure none of the elements is `NaN` or `Inf`. Any advice on how I can track where the loss value is exploding?

**Update**  
I tried retraining much smaller models (K\*4 \approx 80) with the mini batching code above and I’m still getting the exact same error (Loss NaN at second iteration of first epoch). I’m going to investigate further tomorrow but if anyone spotted any obvious errors please do let me know, thanks!

**Update 2**  
Yep, I was an idiot and forgot to do `i+batch_size` in `trainData[:, randIdx[i:batch_size]]`. Now I just need to deal with the `OutOfMemoryError()`.