# Gradient loop vs batch gradient speed

**URL:** <https://discourse.julialang.org/t/gradient-loop-vs-batch-gradient-speed/27209>\
**Category:** Machine Learning\
**Tags:** flux\
**Created:** [August 6, 2019, 8:56am UTC](https://discourse.julialang.org/t/gradient-loop-vs-batch-gradient-speed/27209 "2019-08-06T08:56:42Z")\
**Posts on this page:** 5\
**Page:** 1

<div class="post-metadata">

**Author:** ![CraftyVisage](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/craftyvisage/32/8826_2.png) [@CraftyVisage](https://discourse.julialang.org/u/CraftyVisage)\
**Post date:** [August 6, 2019, 8:56am UTC](https://discourse.julialang.org/t/gradient-loop-vs-batch-gradient-speed/27209/1 "2019-08-06T08:56:42Z")

</div>

I am training a neural network using a manual custom gradient update.

I have a loop that goes through every example and computes the gradient of the loss on that example. I then use that gradient in my custom update. This loop is much slower than evaluating the batch gradient on all the examples at once even though presumably the batch also involves evaluating all the individual gradients. Example code is below.

```julia
import Flux
import Random

function main()
    L = 10
    net = Flux.Chain(Flux.Dense(L, L, Flux.relu),
                     Flux.Dense(L, 1))

    N_examples = 1000
    examples = Random.randn(Float64, L, N_examples)
    labels = Random.randn(Float64, 1, N_examples)

    loss(x, y) = sum((net(x) .- y).^2) / length(y)

    params = Flux.params(net)
    println("Batch gradient is fast")
    @time grad = Flux.Tracker.gradient(() -> loss(examples, labels))

    println("Gradient of each example is slow")
    @time for i in 1:N_examples
        grad = Flux.Tracker.gradient(() -> loss(examples[:, i], labels[i]))
    end
end

main()

```

Example output:

Batch gradient is fast  
0.006367 seconds (632 allocations: 531.516 KiB)  
Gradient of each example is slow  
0.166312 seconds (505.22 k allocations: 15.944 MiB, 3.84% gc time)

Is there any way to make my loop as fast as the batch evaluation?

---

<div class="post-metadata">

**Author:** ![baggepinnen](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/baggepinnen/32/693_2.png) [@baggepinnen](https://discourse.julialang.org/u/baggepinnen)\
**Post date:** [August 6, 2019, 10:04am UTC](https://discourse.julialang.org/t/gradient-loop-vs-batch-gradient-speed/27209/2 "2019-08-06T10:04:28Z")

</div>

You can land somewhere in between by using mini batches

---

<div class="post-metadata">

**Author:** ![baggepinnen](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/baggepinnen/32/693_2.png) [@baggepinnen](https://discourse.julialang.org/u/baggepinnen)\
**Post date:** [August 6, 2019, 10:05am UTC](https://discourse.julialang.org/t/gradient-loop-vs-batch-gradient-speed/27209/3 "2019-08-06T10:05:51Z")

</div>

You can probably not reach the speed of the full batch version since that can exploit very optimized matrix multiply and simd to a greater extent

---

<div class="post-metadata">

**Author:** ![CraftyVisage](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/craftyvisage/32/8826_2.png) [@CraftyVisage](https://discourse.julialang.org/u/CraftyVisage)\
**Post date:** [August 6, 2019, 11:15am UTC](https://discourse.julialang.org/t/gradient-loop-vs-batch-gradient-speed/27209/4 "2019-08-06T11:15:39Z")

</div>

I can’t use mini-batches for what I’m doing. I need the individual gradients. Do you have any suggestions for speeding up the loop itself?

---

<div class="post-metadata">

**Author:** ![baggepinnen](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/baggepinnen/32/693_2.png) [@baggepinnen](https://discourse.julialang.org/u/baggepinnen)\
**Post date:** [August 7, 2019, 12:01am UTC](https://discourse.julialang.org/t/gradient-loop-vs-batch-gradient-speed/27209/5 "2019-08-07T00:01:36Z")

</div>

You can try to reduce the number of allocations due to slicing, but you’ll never escape the overhead of the AD tracker. You can also try using Zygote.jl instead of tracker, check out Flux#zygote branch