Simple L2 regularisation with Flux.params excruciatingly slow

[lf28]

Here is an MWE, in which I try to apply L2 regularisation with a 27 to 27 dense layer; it takes around 3 second to finish 100 iterations without Flux.params but takes 16 mins to run 10 iterations.

pen_l2(x)::Float32 = sum(abs2, x)
model = Chain(Dense(27=>27, bias=false), softmax)
optim = Flux.setup(Flux.Adam(0.01), model) 
@showprogress for epoch in 1:10
    loss, grads = Flux.withgradient(model) do m
        # 0.01f0 * sum(abs2, m[1].weight)
        0.01f0 * sum(pen_l2, Flux.params(m))
    end
    Flux.update!(optim, model, grads[1])
end

Flux version: v0.13.14
Julia: 1.8.5

[ToucheSir]

I'm guessing this is the same issue as #2040. Can you @time each withgradient call and show the output of that?

[lf28]

Hi, sorry for the late reply. Here is the output:

2.267348 seconds (7.88 M allocations: 331.264 MiB, 1.47% gc time, 62.02% compilation time: 1% of which was recompilation)
2.235811 seconds (7.78 M allocations: 326.473 MiB, 1.51% gc time, 62.24% compilation time)
2.260386 seconds (7.78 M allocations: 326.592 MiB, 1.41% gc time, 61.92% compilation time)
2.239555 seconds (7.79 M allocations: 326.710 MiB, 1.31% gc time, 62.14% compilation time)
2.202648 seconds (7.79 M allocations: 326.828 MiB, 1.33% gc time, 62.84% compilation time)
2.201378 seconds (7.79 M allocations: 326.946 MiB, 0.67% gc time, 62.20% compilation time)
2.242019 seconds (7.79 M allocations: 327.064 MiB, 1.31% gc time, 61.98% compilation time)
2.184376 seconds (7.80 M allocations: 327.182 MiB, 1.34% gc time, 63.11% compilation time)
2.189200 seconds (7.80 M allocations: 327.301 MiB, 1.34% gc time, 61.56% compilation time)
2.187360 seconds (7.80 M allocations: 327.419 MiB, 1.34% gc time, 63.10% compilation time)

[ToucheSir]

Looks like one and the same issue. See if #2040 (comment) works for you, and feel free to share more about your model on that issue if it doesn't.

[ToucheSir]

Duplicate of #2040