# Multitask Learning in Flux

**URL:** <https://discourse.julialang.org/t/multitask-learning-in-flux/79535>\
**Category:** Machine Learning\
**Tags:** flux\
**Created:** [April 15, 2022, 1:05pm UTC](https://discourse.julialang.org/t/multitask-learning-in-flux/79535 "2022-04-15T13:05:47Z")\
**Posts on this page:** 3\
**Page:** 1

<div class="post-metadata">

**Author:** ![Marco\_Repetto](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/marco_repetto/32/35451_2.png) [@Marco\_Repetto](https://discourse.julialang.org/u/Marco_Repetto)\
**Post date:** [April 15, 2022, 1:05pm UTC](https://discourse.julialang.org/t/multitask-learning-in-flux/79535/1 "2022-04-15T13:05:47Z")

</div>

I would like to implement an MTL NN in Flux but is not clear to me how to create a different set of layers for each task.

Right now I’m outputting all the tasks in the final layer and using `logitcrossentropy` loss as all the tasks involve classification. The model is working but I would like to apply a different loss for each task as well as custom layers that are task specific.

```julia
n_in = size(train_x)[1]
n_out = size(train_y)[1]

skip_component = SkipConnection(
    Chain(
        Dropout(0.15),
        Dense(n_in, n_in*2, relu),
        BatchNorm(n_in*2, relu),
        Dropout(0.15),
        Dense(n_in*2, n_in*2, relu),
        Dropout(0.15),
        Dense(n_in*2, n_in*2, relu),
        Dropout(0.15),
        Dense(n_in*2, n_in*2, relu),
        BatchNorm(n_in*2, relu),
        Dropout(0.15),
        Dense(n_in*2, n_in, relu),
    ), +)
m = Chain(
    Dropout(0.3),
    Dense(n_in, n_in, relu),
    skip_component,
    skip_component,
    skip_component,
    skip_component,
    skip_component,
    skip_component,
    skip_component,
    skip_component,
    skip_component,
    Dense(n_in, n_out, sigmoid),
) |> gpu

```

---

<div class="post-metadata">

**Author:** ![FPGro](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/fpgro/32/20822_2.png) [@FPGro](https://discourse.julialang.org/u/FPGro)\
**Post date:** [April 15, 2022, 3:36pm UTC](https://discourse.julialang.org/t/multitask-learning-in-flux/79535/2 "2022-04-15T15:36:58Z")

</div>

Hi, did you read  
[https://fluxml.ai/Flux.jl/stable/models/advanced/#Custom-multiple-input-or-output-layer](https://fluxml.ai/Flux.jl/stable/models/advanced/#Custom-multiple-input-or-output-layer),  
especially the latter part about a custom split function?

You should be able to split the shared body into task-specific paths pretty easily this way. And if you modify the last listing on that page a bit, I guss you can use a different loss for each comppnent of the resulting tuple. You just need to combine it somehow to get a single number out.

---

<div class="post-metadata">

**Author:** ![Marco\_Repetto](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/marco_repetto/32/35451_2.png) [@Marco\_Repetto](https://discourse.julialang.org/u/Marco_Repetto)\
**Post date:** [April 15, 2022, 4:42pm UTC](https://discourse.julialang.org/t/multitask-learning-in-flux/79535/3 "2022-04-15T16:42:29Z")

</div>

No I did not! Thank you for pointing this out!