Multitask Learning in Flux

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.

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

Hi, did you read
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.

1 Like

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