How to train a combination of models in Flux?

General Usage
1

I have two models m1 and m2. Here is my code:

using Flux

function even_mask(x)
    s1, s2 = size(x)
    weight_mask = zeros(s1, s2)
    weight_mask[2:2:s1,:] = ones(Int(s1/2), s2)
    return weight_mask
end

function odd_mask(x)
    s1, s2 = size(x)
    weight_mask = zeros(s1, s2)
    weight_mask[1:2:s1,:] = ones(Int(s1/2), s2)
    return weight_mask
end

function even_duplicate(x)
    s1, s2 = size(x)
    x_ = zeros(s1, s2)
    x_[1:2:s1,:] = x[1:2:s1,:]
    x_[2:2:s1,:] = x[1:2:s1,:]
    return x_
end

function odd_duplicate(x)
    s1, s2 = size(x)
    x_ = zeros(s1, s2)
    x_[1:2:s1,:] = x[2:2:s1,:]
    x_[2:2:s1,:] = x[2:2:s1,:]
    return x_
end

function Even(m)
    x -> x .+ even_mask(x).*m(even_duplicate(x))
end

function InvEven(m)
    x -> x .- even_mask(x).*m(even_duplicate(x))
end

function Odd(m)
    x -> x .+ odd_mask(x).*m(odd_duplicate(x))
end

function InvOdd(m)
    x -> x .- odd_mask(x).*m(odd_duplicate(x))
end

m1 = Chain(Dense(4,6,relu), Dense(6,5,relu), Dense(5,4))
m2 = Chain(Dense(4,7,relu), Dense(7,4))

forward = Chain(Even(m1), Odd(m2))
inverse = Chain(InvOdd(m2), InvEven(m1))

function loss(x)
    z = forward(x)
    return 0.5*sum(z.*z)
end

opt = Flux.ADAM()

x = rand(4,100)

for i=1:100
    Flux.train!(loss, Flux.params(forward), x, opt)
    println(loss(x))
end

The forward model is a combination of m1 and m2. I need to optimize m1 and m2 so I could optimize both forward and inverse models. But it seems that params(forward) is empty. How could I train my model?

2

try with passing Flux.params(m1, m2) to train!.
I see you are doing array mutations, that won’t work with zygote

3

What is array mutation?