# How to use Flux.stop()

**URL:** <https://discourse.julialang.org/t/how-to-use-flux-stop/77746>\
**Category:** Machine Learning\
**Tags:** flux\
**Created:** [March 11, 2022, 12:19pm UTC](https://discourse.julialang.org/t/how-to-use-flux-stop/77746 "2022-03-11T12:19:35Z")\
**Posts on this page:** 2\
**Page:** 1

<div class="post-metadata">

**Author:** ![TheLateKronos](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/thelatekronos/32/12824_2.png) [@TheLateKronos](https://discourse.julialang.org/u/TheLateKronos)\
**Post date:** [March 11, 2022, 12:19pm UTC](https://discourse.julialang.org/t/how-to-use-flux-stop/77746/1 "2022-03-11T12:19:36Z")

</div>

I am trying to learn how to use `Flux.stop()` in a callback. My problem is that as far as I can tell, my call to `Flux.stop()` has litterally no effect - see the following example:

The setup:

> ****
>
> ```julia
> using Flux
> x_train = randn(Float32, 28, 28, 1, 10_000)
> labels_train = Flux.onehotbatch(rand(0:9, 10_000), 0:9)
> loader_train = Flux.DataLoader((x_train, labels_train), batchsize=100)
> model = Chain(
> Conv((5, 5), 1=>8, pad=2, stride=2, relu),
> Conv((3, 3), 8=>16, pad=1, stride=2, relu),
> Conv((3, 3), 16=>32, pad=1, stride=2, relu),
> GlobalMeanPool(),
> flatten,
> Dense(32, 10),
> softmax
> ) # 6_346 parameters
> loss = (x, y)->Flux.Losses.logitcrossentropy(model(x), y)
> opt = ADAM()
> parameters = params(model)
> 
> time_limit = 1 # Seconds
> 
> ```

With that setup, lets try looping over the dataloader:

> ****
>
> ```julia
> times = Dict(:initial => time(), :elapsed => time())
> for (xs, labels) in loader_train
> Flux.train!(loss, parameters, [(xs, labels)], opt, cb = function cb()
> global times[:elapsed] = time() - times[:initial]
> "Time elapsed = $(round(times[:elapsed], digits=2))" |> println
> if times[:elapsed] > time_limit
> println("Trained for $(round(times[:elapsed], digits=2)), stopping")
> Flux.stop()
> end
> end
> )
> end
> 
> ```

This produces the following output:

> ****
>
> ```julia
> Time elapsed = 0.12
> Time elapsed = 0.13
> Time elapsed = 0.14
> Time elapsed = 0.15
> Time elapsed = 0.16
> Time elapsed = 0.17
> Time elapsed = 0.18
> Time elapsed = 0.19
> Time elapsed = 0.2
> Time elapsed = 0.21
> Time elapsed = 0.22
> Time elapsed = 0.23
> Time elapsed = 0.24
> Time elapsed = 0.25
> Time elapsed = 0.26
> Time elapsed = 0.27
> Time elapsed = 0.28
> Time elapsed = 0.29
> Time elapsed = 0.3
> Time elapsed = 0.31
> Time elapsed = 0.32
> Time elapsed = 0.33
> Time elapsed = 0.34
> Time elapsed = 0.35
> Time elapsed = 0.36
> Time elapsed = 0.37
> Time elapsed = 0.38
> Time elapsed = 0.38
> Time elapsed = 0.39
> Time elapsed = 0.4
> Time elapsed = 0.41
> Time elapsed = 0.42
> Time elapsed = 0.43
> Time elapsed = 0.44
> Time elapsed = 0.45
> Time elapsed = 0.45
> Time elapsed = 0.46
> Time elapsed = 0.47
> Time elapsed = 0.48
> Time elapsed = 0.49
> Time elapsed = 0.5
> Time elapsed = 0.51
> Time elapsed = 0.52
> Time elapsed = 0.53
> Time elapsed = 0.54
> Time elapsed = 0.55
> Time elapsed = 0.56
> Time elapsed = 0.57
> Time elapsed = 0.58
> Time elapsed = 0.58
> Time elapsed = 0.59
> Time elapsed = 0.6
> Time elapsed = 0.61
> Time elapsed = 0.62
> Time elapsed = 0.63
> Time elapsed = 0.64
> Time elapsed = 0.65
> Time elapsed = 0.66
> Time elapsed = 0.67
> Time elapsed = 0.67
> Time elapsed = 0.68
> Time elapsed = 0.69
> Time elapsed = 0.7
> Time elapsed = 0.71
> Time elapsed = 0.84
> Time elapsed = 0.85
> Time elapsed = 0.86
> Time elapsed = 0.86
> Time elapsed = 0.87
> Time elapsed = 0.88
> Time elapsed = 0.89
> Time elapsed = 0.9
> Time elapsed = 0.91
> Time elapsed = 0.92
> Time elapsed = 0.93
> Time elapsed = 0.94
> Time elapsed = 0.95
> Time elapsed = 0.95
> Time elapsed = 0.96
> Time elapsed = 0.97
> Time elapsed = 0.98
> Time elapsed = 0.99
> Time elapsed = 1.0
> Time elapsed = 1.01
> Trained for 1.01, stopping
> Time elapsed = 1.02
> Trained for 1.02, stopping
> Time elapsed = 1.03
> Trained for 1.03, stopping
> Time elapsed = 1.04
> Trained for 1.04, stopping
> Time elapsed = 1.05
> Trained for 1.05, stopping
> Time elapsed = 1.05
> Trained for 1.05, stopping
> Time elapsed = 1.06
> Trained for 1.06, stopping
> Time elapsed = 1.07
> Trained for 1.07, stopping
> Time elapsed = 1.08
> Trained for 1.08, stopping
> Time elapsed = 1.09
> Trained for 1.09, stopping
> Time elapsed = 1.1
> Trained for 1.1, stopping
> Time elapsed = 1.11
> Trained for 1.11, stopping
> Time elapsed = 1.12
> Trained for 1.12, stopping
> Time elapsed = 1.13
> Trained for 1.13, stopping
> Time elapsed = 1.13
> Trained for 1.13, stopping
> Time elapsed = 1.14
> Trained for 1.14, stopping
> Time elapsed = 1.15
> Trained for 1.15, stopping
> 
> ```

So no stopping occurs - it seems like nothing changes when the elapsed time goes over the timelimit.

So lets try a single pass over all the data:

> ****
>
> ```julia
> times = Dict(:initial => time(), :elapsed => time())
> for _ in 1:3 # Repeat the for loop below 3 times
> for (xs, labels) in [(x_train[:, :, :, 1:10_000], labels_train[:, 1:10_000])]
> t₀ = time()
> Flux.train!(loss, parameters, [(xs, labels)], opt, cb = function cb()
> global times[:elapsed] = time() - times[:initial]
> "Time elapsed = $(round(times[:elapsed], digits=2))" |> println
> if times[:elapsed] > time_limit
> println("Trained for $(round(times[:elapsed], digits=2)), stopping")
> Flux.stop()
> end
> end
> )
> end
> end
> 
> ```

This produces the following output:

> ****
>
> ```julia
> Time elapsed = 0.93
> Time elapsed = 1.82
> Trained for 1.82, stopping
> Time elapsed = 2.6
> Trained for 2.6, stopping
> 
> ```

Again, no stopping occurs - after going over 1 second trained, nothing changes.

So it seems like the callback prints things just fine, but at no point actually stops the training. What is going wrong?

---

<div class="post-metadata">

**Author:** ![trahflow](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/trahflow/32/30585_2.png) [@trahflow](https://discourse.julialang.org/u/trahflow)\
**Post date:** [March 11, 2022, 2:14pm UTC](https://discourse.julialang.org/t/how-to-use-flux-stop/77746/2 "2022-03-11T14:14:54Z")

</div>

I think `Flux.stop()` is not meant to stop any outer loop around `Flux.train!()`.

But for your example you do not actually need a loop around `Flux.train!()`. Instead what you can do is to directly provide your dataloader to `Flux.train!()` (the loop over minibatches happens _inside_ `Flux.train!()`).  
See the minimal example here: [Training · Flux](https://fluxml.ai/Flux.jl/stable/training/training/#Loss-Functions)

Btw., `Flux.train!()` and `Flux.stop()` are no magic. The latter just throws a custom error, that the former catches within the minibatch loop: [https://github.com/FluxML/Flux.jl/blob/ef04fda844ea05c45d8cc53448e3b25513a77617/src/optimise/train.jl#L63-L124](https://github.com/FluxML/Flux.jl/blob/ef04fda844ea05c45d8cc53448e3b25513a77617/src/optimise/train.jl#L63-L124)

So if you replace your training loop (your second code block) by:

```julia
Flux.train!(loss, parameters, loader_train, opt, cb = function cb()
     global times[:elapsed] = time() - times[:initial]
     "Time elapsed = $(round(times[:elapsed], digits=2))" |> println
     if times[:elapsed] > time_limit
         println("Trained for $(round(times[:elapsed], digits=2)), stopping")
         Flux.stop()
     end
 end
 )

```

this produces what you want:

```julia
...
Time elapsed = 0.94
Time elapsed = 0.95
Time elapsed = 0.97
Time elapsed = 1.01
Trained for 1.01, stopping

```