# Flux: combine two neural networks to run simultaneously?

**URL:** <https://discourse.julialang.org/t/flux-combine-two-neural-networks-to-run-simultaneously/120776>\
**Category:** Machine Learning\
**Tags:** flux\
**Created:** [October 1, 2024, 9:22am UTC](https://discourse.julialang.org/t/flux-combine-two-neural-networks-to-run-simultaneously/120776 "2024-10-01T09:22:46Z")\
**Posts on this page:** 5\
**Page:** 1

<div class="post-metadata">

**Author:** ![greatpet](https://avatars.discourse-cdn.com/v4/letter/g/e495f1/32.png) [@greatpet](https://discourse.julialang.org/u/greatpet)\
**Post date:** [October 1, 2024, 9:22am UTC](https://discourse.julialang.org/t/flux-combine-two-neural-networks-to-run-simultaneously/120776/1 "2024-10-01T09:22:47Z")

</div>

For example, I have an input which is a length-10 vector. I want the first neural network to output a length-3 vector based on the first 5 elements of the input, and the second neural network to output a length-4 vector based on the last 5 element of the input. The overall output should be a length-7 vector which combines the outputs of the two neural networks. What’s the easiest way to do this? Though this sounds strange, the external package I use only accepts a “single” neural network in its API, so I cannot just use the two networks separately.

---

<div class="post-metadata">

**Author:** ![mcabbott](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mcabbott/32/6603_2.png) [@mcabbott](https://discourse.julialang.org/u/mcabbott)\
**Post date:** [October 1, 2024, 1:27pm UTC](https://discourse.julialang.org/t/flux-combine-two-neural-networks-to-run-simultaneously/120776/2 "2024-10-01T13:27:28Z")

</div>

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

```julia
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

```

---

<div class="post-metadata">

**Author:** ![greatpet](https://avatars.discourse-cdn.com/v4/letter/g/e495f1/32.png) [@greatpet](https://discourse.julialang.org/u/greatpet)\
**Post date:** [October 1, 2024, 2:58pm UTC](https://discourse.julialang.org/t/flux-combine-two-neural-networks-to-run-simultaneously/120776/3 "2024-10-01T14:58:08Z")

</div>

The only concern would be the amount of memory allocation involved. Can the input vector be split into two allocation-free views to be fed into the two NNs? Can I pre-allocate a single output vector and ask the two NNs to each write to a specific part of the output?

---

<div class="post-metadata">

**Author:** ![mcabbott](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mcabbott/32/6603_2.png) [@mcabbott](https://discourse.julialang.org/u/mcabbott)\
**Post date:** [October 1, 2024, 3:21pm UTC](https://discourse.julialang.org/t/flux-combine-two-neural-networks-to-run-simultaneously/120776/4 "2024-10-01T15:21:29Z")

</div>

Such optimisations are not impossible, but usually any nontrivial `model1` & `model2` will allocate far more than these simple steps anyway.

---

<div class="post-metadata">

**Author:** ![sylvaticus](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/sylvaticus/32/203883_2.png) [@sylvaticus](https://discourse.julialang.org/u/sylvaticus)\
**Post date:** [October 2, 2024, 12:21pm UTC](https://discourse.julialang.org/t/flux-combine-two-neural-networks-to-run-simultaneously/120776/5 "2024-10-02T12:21:34Z")

</div>

Not in Flux, but in the NN module of [BetaML](https://github.com/sylvaticus/BetaML.jl) you can use `ReplicatorLayer` and `GroupedLayer` to obtain a [multi branch neural network](https://sylvaticus.github.io/BetaML.jl/stable/tutorials/Multi-branch%20neural%20network/betaml_tutorial_multibranch_nn.html):

 ![multibranch_nn](https://global.discourse-cdn.com/julialang/original/3X/6/c/6c215b918a7db3f73603fcac47a15742dbc486c6.png)

Perhaps you can use the same approach in Flux…