Flux: combine two neural networks to run simultaneously?

Specific Domains Machine Learning
2

One way is to use Parallel, applying vcat after the two paths, and some function to split the input before:

julia> mysplit(x::AbstractVector) = (x[1:5], x[6:end]);

julia> model1 = Dense(5 => 3);

julia> model2 = Dense(5 => 4, x->x+100);

julia> model = Chain(mysplit, Parallel(vcat, model1, model2))
Chain(
  mysplit,
  Parallel(
    vcat,
    Dense(5 => 3),                      # 18 parameters
    Dense(5 => 4, #5),                  # 24 parameters
  ),
)                   # Total: 4 arrays, 42 parameters, 376 bytes.

julia> model(ones32(10))
7-element Vector{Float32}:
  -1.6361262
   1.4683602
  -0.70261115
 100.29568
 100.38945
 100.91372
 100.26898

julia> mysplit(x::AbstractMatrix) = (x[1:5, :], x[6:end, :]);

julia> model(ones32(10, 3))  # now accepts a batch of inputs
7×3 Matrix{Float32}:
  -1.63613    -1.63613    -1.63613
   1.46836     1.46836     1.46836
  -0.702611   -0.702611   -0.702611
 100.296     100.296     100.296
 100.389     100.389     100.389
 100.914     100.914     100.914
 100.269     100.269     100.269
1 Like