# Significant compile time latency in Flux with a GAN

**URL:** <https://discourse.julialang.org/t/significant-compile-time-latency-in-flux-with-a-gan/68518>\
**Category:** Machine Learning\
**Tags:** compilation, flux\
**Created:** [September 21, 2021, 10:02am UTC](https://discourse.julialang.org/t/significant-compile-time-latency-in-flux-with-a-gan/68518 "2021-09-21T10:02:00Z")\
**Posts on this page:** 9\
**Page:** 2

<div class="post-metadata">

**Author:** ![Alexander-Barth](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/alexander-barth/32/3692_2.png) [@Alexander-Barth](https://discourse.julialang.org/u/Alexander-Barth)\
**Post date:** [November 26, 2021, 8:54pm UTC](https://discourse.julialang.org/t/significant-compile-time-latency-in-flux-with-a-gan/68518/21 "2021-11-26T20:54:49Z")

</div>

Indeed on the CPU I also see this high compile time latency on the first call (more than ~~10~~ 30 minutes, still running) with Julia 1.6.3.

---

<div class="post-metadata">

**Author:** ![goerch](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/goerch/32/29122_2.png) [@goerch](https://discourse.julialang.org/u/goerch)\
**Post date:** [November 26, 2021, 8:59pm UTC](https://discourse.julialang.org/t/significant-compile-time-latency-in-flux-with-a-gan/68518/22 "2021-11-26T20:59:05Z")

</div>

Question of an uninitiated: does the Julia compiler dispatch compilation of specializations to the GPU packages somehow? Ah, OK: I’d imagine GPU packages use for example macros in functions, whose compilation is obviously dispatched by the compiler depending on optimization (specialization) level?

---

<div class="post-metadata">

**Author:** ![Alexander-Barth](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/alexander-barth/32/3692_2.png) [@Alexander-Barth](https://discourse.julialang.org/u/Alexander-Barth)\
**Post date:** [November 26, 2021, 9:02pm UTC](https://discourse.julialang.org/t/significant-compile-time-latency-in-flux-with-a-gan/68518/23 "2021-11-26T21:02:30Z")

</div>

Here is a screenshot of the callgrind profiles (only for the second call to pullback) on julia 1.6.3 which shows that still a lot of time is spent on compilation (in fact almost all):

 ![Bildschirmfoto von 2021-11-10 09-32-46](https://global.discourse-cdn.com/julialang/original/3X/8/4/8487caaf58782541cc346df9e2d998b0952a30c0.png)

Could it be that LLVM tries to do some complex optimizations but which fails due to the deeply nested types used here. Then on the second call, the compiler wants again optimize the code (and takes again a long time).

---

<div class="post-metadata">

**Author:** ![goerch](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/goerch/32/29122_2.png) [@goerch](https://discourse.julialang.org/u/goerch)\
**Post date:** [November 26, 2021, 9:11pm UTC](https://discourse.julialang.org/t/significant-compile-time-latency-in-flux-with-a-gan/68518/24 "2021-11-26T21:11:46Z")

</div>

I think it has more to do with a combinatorial explosion of specializations induced by deeply nested varying tuple types (because using either -O1 or using `@nospecialize` seems to reduce/eliminate the problem). Only question to discuss on [https://github.com/JuliaLang/julia/issues/43206](https://github.com/JuliaLang/julia/issues/43206) is if this optimization works as intended or out of control in these situations…

Edit: I just checked that and have to correct myself: it doesn’t seem to be the number but the complexity of specialization, which is the problem. In both cases I analyzed the critical signature was a couple of 100 kB long, so yes Alexander, you are probably right!

---

<div class="post-metadata">

**Author:** ![ToucheSir](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/touchesir/32/14411_2.png) [@ToucheSir](https://discourse.julialang.org/u/ToucheSir)\
**Post date:** [November 27, 2021, 6:35am UTC](https://discourse.julialang.org/t/significant-compile-time-latency-in-flux-with-a-gan/68518/25 "2021-11-27T06:35:38Z")

</div>

I managed to reduce this to the following:

```julia
using Flux

channels = 4

function resblock(channels)
    return SkipConnection(Chain(
        Conv((3, 3), channels => channels, pad=1),
        Conv((3, 3), channels => channels, pad=1),
    ), +)
end

model = Chain(
    SkipConnection(
        Chain(
            resblock(channels),
            resblock(channels),
            resblock(channels),
            resblock(channels),
            resblock(channels),

            resblock(channels),
            resblock(channels),
            resblock(channels),
            resblock(channels),
            resblock(channels),

            resblock(channels),
            resblock(channels),
            resblock(channels),
            resblock(channels),
            resblock(channels),
        ),
    +),
    AdaptiveMeanPool((1, 1))
)

display(model)
println()
@show typeof(model) 

loss(x) = sum(model(x))

lr_images = randn(Float32, 2, 2, channels, 1)
@time loss(lr_images)
@time loss(lr_images) 
    
loss_grad(x, ps) = gradient(() -> loss(x), ps)
    
ps = Flux.params(model)
@time loss_grad(lr_images, ps)
@time loss_grad(lr_images, ps)

```

Not sure about the second gradient timing, but the 1st one has been going for over 10 minutes with no results. Importantly, removing just one of the `resblock`s results in a far more tractable compilation time (115s).

Do you mind filing an issue about this? I wonder if [https://github.com/FluxML/Zygote.jl/issues/1119](https://github.com/FluxML/Zygote.jl/issues/1119) might be related as well.

---

<div class="post-metadata">

**Author:** ![DrChainsaw](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/drchainsaw/32/8497_2.png) [@DrChainsaw](https://discourse.julialang.org/u/DrChainsaw)\
**Post date:** [November 27, 2021, 11:15am UTC](https://discourse.julialang.org/t/significant-compile-time-latency-in-flux-with-a-gan/68518/26 "2021-11-27T11:15:22Z")

</div>

I took the “varying tuple types” hypothesis for a spin and tried the same model using `NaiveNASflux` which uses a conciously type unstable representation of the model to allow for mutation.

Using the same code as @ToucheSir above but converting it to a `CompGraph` using `ONNXNaiveNASflux` (as there is no convenient built-in for direct `Chain` =\> `CompGraph` translation):

```julia
mio = PipeBuffer();

# One off thing as ONNXNaiveNASflux currently does not support AdapiveMeanPool
(::AdaptiveMeanPool)(pp::ONNXNaiveNASflux.AbstractProbe) = ONNXNaiveNASflux.globalpool(pp, identity, "GlobalAveragePool");

save(mio, model);

model_graph = load(mio);

model_graph(lr_images) == model(lr_images) # true

loss(x) = sum(model_graph(x));

@time loss_grad(lr_images, Flux.params(model_graph))
24.340328 seconds (49.10 M allocations: 2.556 GiB, 5.01% gc time, 99.83% compilation time)
Grads(...)

```

However, the chain example finishes on my computer albeit a bit slower so maybe I’m not comparing apples to apples here:

```julia
loss(x) = sum(model(x))

@time loss_grad(lr_images, Flux.params(model))
 98.507128 seconds (9.56 M allocations: 499.581 MiB, 0.18% gc time, 99.93% compilation time)
Grads(...)

versioninfo()
Julia Version 1.7.0-rc3
Commit 3348de4ea6 (2021-11-15 08:22 UTC)
Platform Info:
  OS: Windows (x86_64-w64-mingw32)
  CPU: Intel(R) Core(TM) i7-5820K CPU @ 3.30GHz
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-12.0.1 (ORCJIT, haswell)
Environment:
  JULIA_DEPOT_PATH = E:/Programs/julia/.julia
  JULIA_EDITOR = code
  JULIA_NUM_THREADS = 6

```

---

<div class="post-metadata">

**Author:** ![Alexander-Barth](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/alexander-barth/32/3692_2.png) [@Alexander-Barth](https://discourse.julialang.org/u/Alexander-Barth)\
**Post date:** [November 27, 2021, 11:58am UTC](https://discourse.julialang.org/t/significant-compile-time-latency-in-flux-with-a-gan/68518/27 "2021-11-27T11:58:01Z")

</div>

Thanks a lot for reducing the code!

> [@ToucheSir](#):
>
> Do you mind filing an issue about this?

Good idea, I just filled it here:

> <https://github.com/FluxML/Zygote.jl/issues/1126>
>
> I am trying to optimize a super-resolution GAN, and I see very significant compi…le time.
> This issue is discussed here:
> https://discourse.julialang.org/t/significant-compile-time-latency-in-flux-with-a-gan/68518
> 
> @ToucheSir was able to reduce this issue to the code below (thanks a lot!).
> 
> I am using Julia 1.6.3 on Linux and these packages:
> 
> \`\`\`
> (@v1.6) pkg\> st Zygote ZygoteRules Flux
> Status \`~/.julia/environments/v1.6/Project.toml\`
> \[587475ba\] Flux v0.12.8 \`https://github.com/FluxML/Flux.jl.git#master\`
> \[e88e6eb3\] Zygote v0.6.30
> \[700de1a5\] ZygoteRules v0.2.2
> \`\`\`
> 
> It appears that this issue is sensitive to the optimization level and to the julia version used.
> 
> with julia 1.6.3 (default optimization):
> 
> \`\`\`
> 2.129568 seconds (6.12 M allocations: 352.760 MiB, 5.21% gc time, 99.96% compilation time)
> 0.000323 seconds (1.80 k allocations: 138.938 KiB)
> 2740.221683 seconds (62.97 M allocations: 3.596 GiB, 0.05% gc time, 0.44% compilation time)
> 0.004320 seconds (12.04 k allocations: 2.761 MiB)
> \`\`\`
> Note the timing of the first call the \`loss\_grad(lr\_images, ps)\`.
> 
> In julia 1.5.3 I get the following timing:
> 
> \`\`\`
> 1.969913 seconds (8.92 M allocations: 458.885 MiB, 6.31% gc time)
> 0.000265 seconds (1.30 k allocations: 102.344 KiB)
> 25.426727 seconds (57.88 M allocations: 2.927 GiB, 3.93% gc time)
> 0.005641 seconds (14.71 k allocations: 1.047 MiB)
> \`\`\`
> 
> with "julia 1.6.3 -O1":
> 
> \`\`\`
> 1.389465 seconds (6.12 M allocations: 352.764 MiB, 25.61% gc time, 99.94% compilation time)
> 0.000403 seconds (2.10 k allocations: 143.625 KiB)
> 28.301099 seconds (62.99 M allocations: 3.596 GiB, 3.96% gc time, 26.40% compilation time)
> 0.004889 seconds (12.94 k allocations: 2.775 MiB)
> \`\`\`
> 
> 
> 
> \`\`\`julia
> using Flux
> 
> channels = 4
> 
> function resblock(channels)
> return SkipConnection(Chain(
> Conv((3, 3), channels =\> channels, pad=1),
> Conv((3, 3), channels =\> channels, pad=1),
> ), +)
> end
> 
> model = Chain(
> SkipConnection(
> Chain(
> resblock(channels),
> resblock(channels),
> resblock(channels),
> resblock(channels),
> resblock(channels),
> 
> resblock(channels),
> resblock(channels),
> resblock(channels),
> resblock(channels),
> resblock(channels),
> 
> resblock(channels),
> resblock(channels),
> resblock(channels),
> resblock(channels),
> resblock(channels),
> ),
> +),
> AdaptiveMeanPool((1, 1))
> )
> 
> display(model)
> println()
> @show typeof(model) 
> 
> loss(x) = sum(model(x))
> 
> lr\_images = randn(Float32, 2, 2, channels, 1)
> @time loss(lr\_images)
> @time loss(lr\_images) 
>     
> loss\_grad(x, ps) = gradient(() -\> loss(x), ps)
>     
> ps = Flux.params(model)
> @time loss\_grad(lr\_images, ps)
> @time loss\_grad(lr\_images, ps)
> \`\`\`
> 
> 
> This issue might be related
> https://github.com/FluxML/Zygote.jl/issues/1119

---

<div class="post-metadata">

**Author:** ![Alexander-Barth](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/alexander-barth/32/3692_2.png) [@Alexander-Barth](https://discourse.julialang.org/u/Alexander-Barth)\
**Post date:** [November 28, 2021, 9:20pm UTC](https://discourse.julialang.org/t/significant-compile-time-latency-in-flux-with-a-gan/68518/28 "2021-11-28T21:20:39Z")

</div>

This looks very interesting! Thank you for pointing me in this direction.

I wanted to compare the gradients. It seems that `length(grad_CompGraph.grads)` has 61 elements while `length(grad_Chain.grads)` has 60 elements.  
But all gradients of the 60 parameters (returned by `Flux.params`) are the same, so I guess the additional elements would not matter.

I am wondering about this line:

```julia
(::AdaptiveMeanPool)(pp::ONNXNaiveNASflux.AbstractProbe) = ONNXNaiveNASflux.globalpool(pp, identity, "GlobalAveragePool");

```

It seems that this actually implements the `AdaptiveMeanPool` in `ONNXNaiveNASflux` (or is it rather a stub)?

(I used `ONNXNaiveNASflux` in the past to convert a ONNX to model Flux, and it worked really great ! 😃 )

```julia
using ONNXNaiveNASflux
using Test

mio = PipeBuffer();

# One off thing as ONNXNaiveNASflux currently does not support AdapiveMeanPool                                                                                        
(::AdaptiveMeanPool)(pp::ONNXNaiveNASflux.AbstractProbe) = ONNXNaiveNASflux.globalpool(pp, identity, "GlobalAveragePool");

save(mio, model);

model_graph = load(mio);

model_graph(lr_images) == model(lr_images) # true                                                                                                                     

loss(x) = sum(model_graph(x));

grad_CompGraph = @time loss_grad(lr_images, Flux.params(model_graph))

loss(x) = sum(model(x))

grad_Chain = @time loss_grad(lr_images, Flux.params(model))

@show length(grad_CompGraph.grads)
@show length(grad_Chain.grads)

for i = 1:length(Flux.params(model))
    @show i
    @show @test grad_Chain[Flux.params(model)[i]] ≈ grad_CompGraph[Flux.params(model_graph)[i]]
end

```

---

<div class="post-metadata">

**Author:** ![DrChainsaw](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/drchainsaw/32/8497_2.png) [@DrChainsaw](https://discourse.julialang.org/u/DrChainsaw)\
**Post date:** [November 29, 2021, 12:00am UTC](https://discourse.julialang.org/t/significant-compile-time-latency-in-flux-with-a-gan/68518/29 "2021-11-29T00:00:32Z")

</div>

> [@Alexander-Barth](#):
>
> It seems that `length(grad_CompGraph.grads)` has 61 elements while `length(grad_Chain.grads)` has 60 elements.

That’s odd. I could swear it was the same as chain when I tested. I’m not at the computer now but I will check again when I am.

> [@Alexander-Barth](#):
>
> It seems that this actually implements the `AdaptiveMeanPool` in `ONNXNaiveNASflux` (or is it rather a stub)?

It adds it as a global pool which happens to be the same in this case (I think, didn’t bother to check how arguments are used), but not in general. If you need it it should not be hard to add assuming the ONNX definition is somewhat sane.

[Previous page](https://discourse.julialang.org/t/significant-compile-time-latency-in-flux-with-a-gan/68518.md?page=1)