# M. learning with regularization using Flux is too slow?

**URL:** <https://discourse.julialang.org/t/m-learning-with-regularization-using-flux-is-too-slow/105683>\
**Category:** Performance\
**Tags:** question, flux\
**Created:** [November 2, 2023, 9:21am UTC](https://discourse.julialang.org/t/m-learning-with-regularization-using-flux-is-too-slow/105683 "2023-11-02T09:21:31Z")\
**Posts on this page:** 6\
**Page:** 1

<div class="post-metadata">

**Author:** ![taka255](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/taka255/32/51426_2.png) [@taka255](https://discourse.julialang.org/u/taka255)\
**Post date:** [November 2, 2023, 9:21am UTC](https://discourse.julialang.org/t/m-learning-with-regularization-using-flux-is-too-slow/105683/1 "2023-11-02T09:21:31Z")

</div>

I am trying to do supervised learning with Flux, as a way to add a regularization term to the loss function referring to the [official documentation](https://fluxml.ai/Flux.jl/stable/training/training/#Regularisation).

I have tried two ways to add a regularization term to the loss function, but adding a penalty directly to the loss function does not finish, whereas using WeightDecay(0.001) the calculation takes only ~2sec. Why does this happen?

Here is the fast code:

```julia
using Flux, Random, Statistics, Distributions, LinearAlgebra, Plots, ProgressMeter

function data_generation(;N = 1000)
    μ1, μ2 = [2.0f0, 0.0f0], [-2.0f0, 0.0f0]  
    σ1, σ2 = 0.5f0, 0.5f0  

    data1 = rand(MvNormal(μ1, σ1*I), N)
    data2 = rand(MvNormal(μ2, σ2*I), N)
    labels1 = fill(0, N)
    labels2 = fill(1, N)

    plt = scatter(data1[1, :], data1[2, :], label="class1")
    scatter!(plt, data2[1, :], data2[2, :], label="class2")
    display(plt)

    # shuffle
    data = hcat(data1, data2)
    labels = hcat(labels1, labels2)
    idxs = shuffle(1:2N)
    data = data[:, idxs]
    labels = labels[idxs]

    # split data
    train_data = [(data[:, i], labels[i]) for i in 1:300]
    test_data = [(data[:, i], labels[i]) for i in 301:N]

    return train_data, test_data
end

# validation
function accuracy(data, model) 
    correct = 0
    total = 0
    for (x, y) in data
        pred = model(x)[1]
        if round(Int64, pred) == y
            correct += 1
        end
        total += 1
    end

    return correct / total
end

function train(;epochs = 1000)
    # data
    train_data, test_data = data_generation() 

    # model and loss
    model = Chain(Dense(2, 1, σ, init=Flux.glorot_normal), sigmoid) 
    loss(y_hat, y) = sum((y_hat .- y).^2) 
    println("Initial accuracy: ", accuracy(train_data, model))
    println("Initial accuracy: ", accuracy(test_data, model))

    # training
    train_accuracy = zeros(epochs)
    test_accuracy = zeros(epochs)
    opt = Flux.setup(OptimiserChain(WeightDecay(0.001), Adam()), model) ## fast
    # opt = Flux.setup(Adam(), model) ## too slow

    pen_l2(x::AbstractArray) = sum(abs2, x)/2
    @showprogress for epoch in 1:epochs
        for data in train_data
            input, label = data
            grads = Flux.gradient(model) do m
                result = m(input)
                #penalty = sum(pen_l2, Flux.params(m))
                loss(result, label) #+ 0.001f0 * penalty
            end

            Flux.update!(opt, model, grads[1])
        end
        train_accuracy[epoch] = accuracy(train_data, model)
        test_accuracy[epoch] = accuracy(test_data, model)
    end

    return train_accuracy, test_accuracy
end

@time train_accuracy, test_accuracy = train()

```

And here is the slow code:

```julia
using Flux, Random, Statistics, Distributions, LinearAlgebra, Plots, ProgressMeter

function train(;epochs = 1000)
    # data
    train_data, test_data = data_generation() 

    # model and loss
    model = Chain(Dense(2, 1, σ, init=Flux.glorot_normal), sigmoid) 
    loss(y_hat, y) = sum((y_hat .- y).^2) 
    println("Initial accuracy: ", accuracy(train_data, model))
    println("Initial accuracy: ", accuracy(test_data, model))

    # training
    train_accuracy = zeros(epochs)
    test_accuracy = zeros(epochs)
    #opt = Flux.setup(OptimiserChain(WeightDecay(0.001), Adam()), model) ## fast
    opt = Flux.setup(Adam(), model) ## too slow

    pen_l2(x::AbstractArray) = sum(abs2, x)/2
    @showprogress for epoch in 1:epochs
        for data in train_data
            input, label = data
            grads = Flux.gradient(model) do m
                result = m(input)
                penalty = sum(pen_l2, Flux.params(m))
                loss(result, label) + 0.001f0 * penalty
            end

            Flux.update!(opt, model, grads[1])
        end
        train_accuracy[epoch] = accuracy(train_data, model)
        test_accuracy[epoch] = accuracy(test_data, model)
    end

    return train_accuracy, test_accuracy
end

@time train_accuracy, test_accuracy = train()

```

Thank you in advance.

---

<div class="post-metadata">

**Author:** ![dmetivie](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/dmetivie/32/6926_2.png) [@dmetivie](https://discourse.julialang.org/u/dmetivie)\
**Post date:** [February 8, 2024, 8:49am UTC](https://discourse.julialang.org/t/m-learning-with-regularization-using-flux-is-too-slow/105683/2 "2024-02-08T08:49:34Z")

</div>

I have the same question!  
This is relevant especially when the regularization is not on all the parameters, but only for specific layer, hence the `WeightDecay` approach in the optimizer is not adapted (?).

It seems that `Flux.params(m)` is taking very long. Does anyone have a better way to regularized?

---

<div class="post-metadata">

**Author:** ![darsnack](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/darsnack/32/10144_2.png) [@darsnack](https://discourse.julialang.org/u/darsnack)\
**Post date:** [February 8, 2024, 3:04pm UTC](https://discourse.julialang.org/t/m-learning-with-regularization-using-flux-is-too-slow/105683/3 "2024-02-08T15:04:22Z")

</div>

The easiest solution here is to use the code in [RFC: add `total` by mcabbott · Pull Request #57 · FluxML/Optimisers.jl · GitHub](https://github.com/FluxML/Optimisers.jl/pull/57) for a faster way to sum parameters. Layer specific regularization can be done by calling `total` on each piece and adding the results together.

---

<div class="post-metadata">

**Author:** ![dmetivie](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/dmetivie/32/6926_2.png) [@dmetivie](https://discourse.julialang.org/u/dmetivie)\
**Post date:** [February 9, 2024, 10:40am UTC](https://discourse.julialang.org/t/m-learning-with-regularization-using-flux-is-too-slow/105683/4 "2024-02-09T10:40:59Z")

</div>

Thanks for the ref.

If I just want to apply regularization to weigths of specific layers, can I just use

```julia
weights_regularizer = lw * Flux.Functors.fmap(pen_l2, layer.weight) # kind of similar to what is done in total
# or even simpler but is it correct?
weights_regularizer = lw * pen_l2(layer.weight) 

```

---

<div class="post-metadata">

**Author:** ![dmetivie](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/dmetivie/32/6926_2.png) [@dmetivie](https://discourse.julialang.org/u/dmetivie)\
**Post date:** [February 9, 2024, 11:15am UTC](https://discourse.julialang.org/t/m-learning-with-regularization-using-flux-is-too-slow/105683/6 "2024-02-09T11:15:30Z")

</div>

I think this is what I was [looking for](https://fluxml.ai/Flux.jl/stable/destructure/#All-Layers) no need to call Param or fmap and it is fast.

However, the `Flux.modules` call is quite long in my example! (but not quite as `Params`)

Ho no! In fact

```julia
ERROR: Flux.modules is not at present differentiable, sorry

```

So why is it in the doc as a solution for regularization? That is a bit confusing.

---

<div class="post-metadata">

**Author:** ![darsnack](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/darsnack/32/10144_2.png) [@darsnack](https://discourse.julialang.org/u/darsnack)\
**Post date:** [February 9, 2024, 12:32pm UTC](https://discourse.julialang.org/t/m-learning-with-regularization-using-flux-is-too-slow/105683/7 "2024-02-09T12:32:49Z")

</div>

`Flux.modules` was created in the days when implicit parameters (`Flux.params`) was the recommended way to take gradients. It looks like in the switch over the explicit parameters, the `Flux.modules` was never properly updated. It should be deprecated and removed from the docs.

The second option you mention can work as long as `layer` refers to the argument being differentiated.

```julia
function loss(m, x, y)
     l = # some loss term
     p = pen_l2(m[1].weight)

    return l + lw * p

```

Notice that I using `m` to get the weight.