Multi Input Multi Output CNN with Flux

Specific Domains Machine Learning
,
1

Hi,

I am super new to machine learning and the Flux.jl package and I am having trouble setting up my model. I am effectively trying to train a dataset of 3 images to predict an output of 2 different images. The inputs are of size [1026,514,3,1] and the outputs are of size [1026,514,2,1]. To do this I have been trying to use Join and Split from Fluxperimental.jl but I am getting method errors with my implementation. Am I implementing this correctly? I am also unsure if I am using Upsample correctly but it hasn’t caused any direct issues for me yet. Any advice is sincerely appreciated.

function m3i2o(max_filters=64, k=5, pad=SamePad(), activation_func=relu, weights_init=Flux.glorot_uniform, mem=gpu)
    function input_layer()
        Chain(
            #downsampling
            Conv((k,k), 3 => Int(round(max_filters/4)), pad=pad, activation_func; init=weights_init),
            MaxPool((2,2)),
            Conv((k,k), Int(round(max_filters/4)) => Int(round(max_filters/2)), pad=pad, activation_func; init=weights_init),
            MaxPool((2,2)),
            Conv((k,k), Int(round(max_filters/2)) => Int(round(max_filters)), pad=pad, activation_func; init=weights_init),
            Conv((k,k), Int(round(max_filters)) => Int(round(max_filters)), pad=pad, activation_func; init=weights_init),
            
            #upsampling
            Upsample(:nearest, scale=(2,2)),
            Conv((k,k), Int(round(max_filters)) => Int(round(max_filters/2)), pad=pad, activation_func; init=weights_init),
            Upsample(:nearest, scale=(2,2)),
            Conv((k,k), Int(round(max_filters/2)) => Int(round(max_filters/4)), pad=pad, activation_func; init=weights_init),
            Conv((k,k), Int(round(max_filters/4)) => 3, pad=pad, activation_func; init=weights_init),
        )
    end

    X = Join(vcat, input_layer())

    model = Chain(
        Split(X, Conv((1,1), 1 => 2, pad=0, activation_func))
    ) |> mem
    return model
end