# Flux output layer with custom activation function

**URL:** <https://discourse.julialang.org/t/flux-output-layer-with-custom-activation-function/98821>\
**Category:** Machine Learning\
**Tags:** question, gpu, flux, type-stability\
**Created:** [May 14, 2023, 3:58am UTC](https://discourse.julialang.org/t/flux-output-layer-with-custom-activation-function/98821 "2023-05-14T03:58:44Z")\
**Posts on this page:** 3\
**Page:** 1

<div class="post-metadata">

**Author:** ![YuriS](https://avatars.discourse-cdn.com/v4/letter/y/ccd318/32.png) [@YuriS](https://discourse.julialang.org/u/YuriS)\
**Post date:** [May 14, 2023, 3:58am UTC](https://discourse.julialang.org/t/flux-output-layer-with-custom-activation-function/98821/1 "2023-05-14T03:58:45Z")

</div>

Hello!  
I am trying to create a neural network that is mostly standard, except with an output layer that applies a sigmoid activation function to one of its four outputs, and just identity to the other three outputs.  
I made a custom activation function following [this thread](https://discourse.julialang.org/t/how-to-apply-an-activation-function-to-a-subset-of-output-units/33202/4):

```julia
custom_final_activation(x::AbstractArray{Float32,2}) = vcat(
    NNlib.sigmoid(x[1:1,:]), NNlib.identity(x[2:end,:]) # PROBLEMATIC?
)

```

and included this as my final layer to my NN model (I have the full, (not)-working minimal example at the end of this post):

```julia
# append final layer to on-going list of layers
push!(layers, Dense(n_intermediate, n_OUT, custom_final_activation))
# construct model
model = Chain(
    layers...
) |> gpu;

```

However, when I try to evaluate this model, I get the error

```julia
ERROR: LoadError: GPU broadcast resulted in non-concrete element type Union{}.
This probably means that the function you are broadcasting contains an error or type instability.

```

I suspect that I have defined my `custom_final_activation` function incorrectly, but I am not sure where the type instability is coming from…  
I would appreciate any help!

Here is a (not)-working minimal example:

```julia
using Flux
using CUDA

n_IN = 10
n_OUT = 4 
n_intermediate = 128
intermediate_layers = 3

# example data 
xs = rand(Float32, n_IN, 100)
ys = rand(Float32, n_OUT, 100)

# build model
layers = Any[Dense(n_IN, n_intermediate, NNlib.relu)]
for i = 1:intermediate_layers
    push!(layers, Dense(n_intermediate, n_intermediate, NNlib.relu))
end
# final layer 
custom_final_activation(x::AbstractArray{Float32,2}) = vcat(
    NNlib.sigmoid(x[1:1,:]), NNlib.identity(x[2:end,:]) # PROBLEMATIC?
)
push!(layers, Dense(n_intermediate, n_OUT, custom_final_activation))
# chain layers together
model = Chain(
    layers...
) |> gpu;

# pull out data 
const bs = 20
train_loader = Flux.DataLoader((xs, ys), batchsize=bs, shuffle = true);
batch = gpu(first(train_loader))

# evaluate model
val, grads = Flux.withgradient(model) do m 
    result = m(batch[1])
    Flux.Losses.mse(result, batch[2])
end

```

---

<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:** [May 14, 2023, 4:05am UTC](https://discourse.julialang.org/t/flux-output-layer-with-custom-activation-function/98821/2 "2023-05-14T04:05:00Z")

</div>

You want something like `Chain(Dense(n_intermediate => n_OUT), custom_final_activation)`.

Activation functions inside Dense are broadcasted, i.e. applied to each number, not to the whole array.

Without the gpu you get a clearer error:

```julia
julia> Dense(2 => 4, custom_final_activation)(rand32(2, 3))
ERROR: MethodError: no method matching custom_final_activation(::Float32)

Closest candidates are:
  custom_final_activation(::AbstractMatrix{Float32})

```

---

<div class="post-metadata">

**Author:** ![YuriS](https://avatars.discourse-cdn.com/v4/letter/y/ccd318/32.png) [@YuriS](https://discourse.julialang.org/u/YuriS)\
**Post date:** [May 14, 2023, 2:31pm UTC](https://discourse.julialang.org/t/flux-output-layer-with-custom-activation-function/98821/3 "2023-05-14T14:31:35Z")

</div>

Thank you, this worked!