# Differentiating the Custom Layer and Regularizer

**URL:** <https://discourse.julialang.org/t/differentiating-the-custom-layer-and-regularizer/82594>\
**Category:** New to Julia\
**Tags:** flux\
**Created:** [June 11, 2022, 9:00am UTC](https://discourse.julialang.org/t/differentiating-the-custom-layer-and-regularizer/82594 "2022-06-11T09:00:07Z")\
**Posts on this page:** 8\
**Page:** 1

<div class="post-metadata">

**Author:** ![ekiral](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/ekiral/32/28879_2.png) [@ekiral](https://discourse.julialang.org/u/ekiral)\
**Post date:** [June 11, 2022, 9:00am UTC](https://discourse.julialang.org/t/differentiating-the-custom-layer-and-regularizer/82594/1 "2022-06-11T09:00:07Z")

</div>

I don’t understand the error `Mutating arrays is not supported`. I couldn’t understand the previous posts on this topic either.

I’m trying to add a very standard weight decay regulariser `R(model)` to my model simply defined as  
`sum(sqnorm, Flux.params(model))`  
where `sqnorm(x) = sum(abs2, x)` just as the tutorial webpages define it.

And in doing `gradient( () -> R(model), Flux.params(model))` I get the annoying message above.

But it doesn’t end there!

I created a custom layer as a struct with positives and negative weights and biases. I get the same message when I apply  
` gradient(() -> Flux.Losses.logitcrossentropy(mapslices(model, x_train, dims = 1), y_train), Flux.params(model))`  
I get the same error message. Actually this also happens when I create a vanilla chained Dense layers. So my custom layer may also contain issues, but the problem seems to be even more basic.

I would very much appreciate if somebody could explain what simple dumb thing I’m doing here. (probably also affecting the performance, is it normal that a full batch loss evaluation takes about 5 seconds on MNIST for a neural network with 28x28, 600, 200, 10 layers?)

---

<div class="post-metadata">

**Author:** ![mcabbott](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mcabbott/32/6603_2.png) [@mcabbott](https://discourse.julialang.org/u/mcabbott)\
**Post date:** [June 11, 2022, 11:03am UTC](https://discourse.julialang.org/t/differentiating-the-custom-layer-and-regularizer/82594/2 "2022-06-11T11:03:40Z")

</div>

> [@ekiral](#):
>
> `sum(sqnorm, Flux.params(model))`  
> where `sqnorm(x) = sum(abs2, x)` just as the tutorial webpages define it.
> 
> And in doing `gradient( () -> R(model), Flux.params(model))` I get the annoying message above.

This ought to work:

```julia
julia> using Flux

julia> sqnorm(x) = sum(abs2, x);

julia> R(model) = sum(sqnorm, Flux.params(model));

julia> model = Chain(Dense(2 => 2, tanh), Dense(2 => 2));

julia> gradient(R, model)
((layers = ((weight = Float32[-1.114341 -0.9922521; 1.7604426 1.4203818], bias = Float32[0.0, 0.0], σ = nothing), (weight = Float32[-2.3257902 -0.59381634; -0.35862422 0.3657079], bias = Float32[0.0, 0.0], σ = nothing)),),)

julia> gr = gradient(() -> R(model), Flux.params(model))
Grads(...)

julia> gr[model[1].weight] # the same as explicit version
2×2 Matrix{Float32}:
 -1.11434 -0.992252
  1.76044 1.42038

```

> [@ekiral](#):
>
> mapslices

This won’t work right now, sadly. You can try [SliceMap](https://github.com/mcabbott/SliceMap.jl); something like that could now be built in but nobody got around to it.

---

<div class="post-metadata">

**Author:** ![ekiral](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/ekiral/32/28879_2.png) [@ekiral](https://discourse.julialang.org/u/ekiral)\
**Post date:** [June 11, 2022, 6:13pm UTC](https://discourse.julialang.org/t/differentiating-the-custom-layer-and-regularizer/82594/3 "2022-06-11T18:13:58Z")

</div>

Thank you very much.

For the regularizer it worked but only after updating the version of Julia to 1.7.3 from 1.5.sth.

And also thank you on using `SliceMap.mapcols` functionality. This works.

But now I have a followup question. What I do isn’t crazy, it should be standard. Do people really need some extra packages outside of Flux in order to calculate loss functions and their gradients? What is the standard way of applying the model to the data points (which are given as columns of a matrix).

I wouldn’t have asked if it weren’t for the fact that calculating the gradient on my medium sized neural network with full batch MNIST took about ~ 50 seconds. That seems to be too high. It was also similar when I used a vanilla dense 3 layer network, and not my custom network.

Also may I ask what was the problem with the `Mutating arrays is not supported` error message? Maybe it will be too hard for me to understand, but for posterity’s sake, could you explain? What in `mapslices` made that error appear?

---

<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:** [June 11, 2022, 9:06pm UTC](https://discourse.julialang.org/t/differentiating-the-custom-layer-and-regularizer/82594/4 "2022-06-11T21:06:43Z")

</div>

> [@ekiral](#):
>
> What I do isn’t crazy, it should be standard.

If you’re referring to including a regularization penalty in the loss, it is standard and @mcabbott’s snippet above shows that it works. Note also [`WeightDecay`](https://fluxml.ai/Flux.jl/stable/training/optimisers/#Flux.Optimise.WeightDecay) which works much as it does in other frameworks: a more efficient way to add L2 regularization to a network.

> [@ekiral](#):
>
> Do people really need some extra packages outside of Flux in order to calculate loss functions and their gradients? What is the standard way of applying the model to the data points (which are given as columns of a matrix).

No. I think the confusion comes from this:

> [@ekiral](#):
>
> gradient(() → Flux.Losses.logitcrossentropy(mapslices(model, x\_train, dims = 1), y\_train), Flux.params(model))

I presume you’re familiar with Python frameworks, so this is basically equivalent to:

```python
criterion = nn.CrossEntropyLoss()
outputs = []
for xs in x_train:
  outputs.append(model(xs))
loss = criterion(outputs, y_train)
loss.backward()

```

Which likewise does not run. Much like you’d call `.backward` once per batch in PyTorch, `gradient` should be called once per batch and not across all batches simultaneously. We try to cover this in [Fitting a Line · Flux](https://fluxml.ai/Flux.jl/stable/models/overview/), so please let us know on the issue tracker if any part of that is unclear.

This also likely explains:

> [@ekiral](#):
>
> I wouldn’t have asked if it weren’t for the fact that calculating the gradient on my medium sized neural network with full batch MNIST took about ~ 50 seconds. That seems to be too high. It was also similar when I used a vanilla dense 3 layer network, and not my custom network.

The other part is that the Julia JIT needs some warmup time to compile everything for the first run of a model. Assuming the loop above is tweaked to work properly with batch(es) as I described earlier, the first run time should be lower and any subsequent runs should be pretty much (sometimes slightly faster/slower) what you’d expect from a Python framework.

> [@ekiral](#):
>
> Also may I ask what was the problem with the `Mutating arrays is not supported` error message? Maybe it will be too hard for me to understand, but for posterity’s sake, could you explain?

For posterity, the [ChainRules docs](https://juliadiff.org/ChainRulesCore.jl/stable/rule_author/which_functions_need_rules.html#Functions-which-mutate-arrays) have a nice section on this topic.

If you look at the part of the error message after the “not supported”, it should tell you _which_ function Zygote is complaining about. Most of the time, that is `setindex!` (i.e. `array[i] = ...`), `push!`/`pop!` or `copyto!` (which comes up in `.=` broadcasting).

If you then look at the stacktrace, you can roughly see where in the call stack is calling a function that does array mutation. In this case, we know it’s somewhere in the implementation of `mapslices`. Now, `mapslices` does not mutate any of its inputs, but there isn’t an explicit AD rule for Zygote to use on it, Zygote will try to drill down as far as it can until it finds a function with an AD rule. This usually works fine for most functions (it’s the “automatic” part of automatic differentiation), but if things bottom out in one of the mutating functions mentioned previously then you’ll receive a mutating arrays error.

---

<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:** [June 11, 2022, 9:21pm UTC](https://discourse.julialang.org/t/differentiating-the-custom-layer-and-regularizer/82594/5 "2022-06-11T21:21:19Z")

</div>

Just to remove any ambiguity about how to run a network over a batch (full or mini) of data in Flux (and most other Julia DL libraries):

```julia
m = Chain(...)
# dummy input, note that the dimensions are *reversed* 
# compared to what you'd use for Python libraries
x = ones(features, batch_size) 
outputs = m(x)
# both outputs and targets have shape (output_features = labels) x batch_size
loss = logitcrossentropy(outputs, targets)

```

So other than the order of dimensions, everything runs fully vectorized and looks very similar to the equivalent Python.

---

<div class="post-metadata">

**Author:** ![ekiral](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/ekiral/32/28879_2.png) [@ekiral](https://discourse.julialang.org/u/ekiral)\
**Post date:** [June 12, 2022, 6:21am UTC](https://discourse.julialang.org/t/differentiating-the-custom-layer-and-regularizer/82594/6 "2022-06-12T06:21:18Z")

</div>

Thank you very much.

Let me maybe also mention why I bothered with mapslices and the like so that other people don’t do the same mistake.

The problem was with my custom layer. When I input `model(x_train)` for x of size `784 x 60000` it didn’t work. The dimensions didn’t match.

Now I checked the definition of Dense layer from the source and saw the mistake of my ways.

I made my layer so that the dense layer could only eat a vector, because as a function it added the bias as a vector. However I should have done `.+` thus broadcasting the bias if necessary and hence allowing for matrices `x` as inputs.

It also had the added benefit of naturally having the performance gain once I gpu’d everything (which was unexpectedly easy to do, though one hickup was that DataLoader outputs normal arrays even if it is fed by CUArrays.)

Once again, thank you.

---

<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:** [June 12, 2022, 7:21pm UTC](https://discourse.julialang.org/t/differentiating-the-custom-layer-and-regularizer/82594/7 "2022-06-12T19:21:40Z")

</div>

Glad to hear things worked out!

> [@ekiral](#):
>
> one hickup was that DataLoader outputs normal arrays even if it is fed by CUArrays.)

This should not be happening, do you have a MWE?

---

<div class="post-metadata">

**Author:** ![ekiral](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/ekiral/32/28879_2.png) [@ekiral](https://discourse.julialang.org/u/ekiral)\
**Post date:** [June 14, 2022, 1:35am UTC](https://discourse.julialang.org/t/differentiating-the-custom-layer-and-regularizer/82594/8 "2022-06-14T01:35:32Z")

</div>

You’re right. It didn’t happen after I closed and ran the Julia code again. My mistake, must have still referenced a cpu version of the DataLoader in the memory.