# Collect model outputs into multi-dimensional array

**URL:** <https://discourse.julialang.org/t/collect-model-outputs-into-multi-dimensional-array/72977>\
**Category:** Machine Learning\
**Tags:** flux\
**Created:** [December 12, 2021, 2:14pm UTC](https://discourse.julialang.org/t/collect-model-outputs-into-multi-dimensional-array/72977 "2021-12-12T14:14:50Z")\
**Posts on this page:** 6\
**Page:** 1

<div class="post-metadata">

**Author:** ![Tengrath](https://avatars.discourse-cdn.com/v4/letter/t/c77e96/32.png) [@Tengrath](https://discourse.julialang.org/u/Tengrath)\
**Post date:** [December 12, 2021, 2:14pm UTC](https://discourse.julialang.org/t/collect-model-outputs-into-multi-dimensional-array/72977/1 "2021-12-12T14:14:50Z")

</div>

Consider the following model which creates an output after every layer

’  
function track(model, input)  
track::Vector{Vector{Float64}} =

```
p = input
for layer in model
    p = layer(p)
    track = vcat(track, [p,])
end
return track

```

end

’

which is unfortunately quite inefficient. Is there a way to store the outputs more efficiently into preallocated memory? Ideally into a mutli-dimensional array?

---

<div class="post-metadata">

**Author:** ![trahflow](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/trahflow/32/30585_2.png) [@trahflow](https://discourse.julialang.org/u/trahflow)\
**Post date:** [December 13, 2021, 1:45pm UTC](https://discourse.julialang.org/t/collect-model-outputs-into-multi-dimensional-array/72977/2 "2021-12-13T13:45:22Z")

</div>

> [@Tengrath](#):
>
> Is there a way to store the outputs more efficiently into preallocated memory? Ideally into a mutli-dimensional array?

Sure, if you know the size(s) beforehand (and things are rectangular, i.e. size of the output of each layer is the same) E.g. for 5 layers, each with an output of length 64, you can initialize a multi(i.e. 2)-dimensional array like this:

```julia
track = Matrix{Float64}(undef, 64, 5)

```

And then loop over the layers like this:

```julia
for (i, layer) in enumerate(model)
    p = layer(p)
    track[:, i] = p
end

```

---

<div class="post-metadata">

**Author:** ![stevengj](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/stevengj/32/71_2.png) [@stevengj](https://discourse.julialang.org/u/stevengj)\
**Post date:** [December 13, 2021, 2:01pm UTC](https://discourse.julialang.org/t/collect-model-outputs-into-multi-dimensional-array/72977/3 "2021-12-13T14:01:32Z")

</div>

> [@trahflow](#):
>
> ```julia
> p = layer(p)
> track[:, i] = p
> 
> ```

Not that this still allocates a vector `p` and then copies it into your matrix. To avoid allocations completely, you would want to rewrite your `layer` function to operate in-place on a pre-allocated vector (and then pass in e.g. `@view track[:, i]` to write the results directly into the matrix column).

---

<div class="post-metadata">

**Author:** ![Tengrath](https://avatars.discourse-cdn.com/v4/letter/t/c77e96/32.png) [@Tengrath](https://discourse.julialang.org/u/Tengrath)\
**Post date:** [December 13, 2021, 6:01pm UTC](https://discourse.julialang.org/t/collect-model-outputs-into-multi-dimensional-array/72977/4 "2021-12-13T18:01:57Z")

</div>

Unfortunately this approach is not supported by Flux/Zygote

> Mutating arrays is not supported – called setindex!(::Matrix{Float64}, \_…)

But I guess there must be some workaround for this since the array is not mutated in a mathematical sense. In contrast to in-place operations which override data necessary for calculation of gradients, this just stores results into memory.

---

<div class="post-metadata">

**Author:** ![ToucheSir](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/touchesir/32/14411_2.png) [@ToucheSir](https://discourse.julialang.org/u/ToucheSir)\
**Post date:** [December 13, 2021, 6:10pm UTC](https://discourse.julialang.org/t/collect-model-outputs-into-multi-dimensional-array/72977/5 "2021-12-13T18:10:37Z")

</div>

I don’t think it was clear up front that this was in a gradient context, because that changes the solution space quite dramatically. There are two ways you can go about this in Zygote:

1. Use [`Zygote.Buffer`](https://fluxml.ai/Zygote.jl/latest/utils/#Zygote.Buffer) instead of a plain array.
2. Put your layers in a `Chain`, call [`Flux.activations`](https://fluxml.ai/Flux.jl/stable/models/regularisation/#Flux.activations) to get a set of outputs and then `reduce(hcat, outputs)` to allocate the array once.

#1 is your best shot (short of writing AD rules) for the in-place option @stevengj described. #2 may be faster if you can live with the `p = layer(p)` allocation for each layer, but I would try both just to be sure.

---

<div class="post-metadata">

**Author:** ![Tengrath](https://avatars.discourse-cdn.com/v4/letter/t/c77e96/32.png) [@Tengrath](https://discourse.julialang.org/u/Tengrath)\
**Post date:** [December 16, 2021, 10:25am UTC](https://discourse.julialang.org/t/collect-model-outputs-into-multi-dimensional-array/72977/6 "2021-12-16T10:25:42Z")

</div>

Thanks a lot!

In case I stick with ‘p = layer(p)’ [Flux.activations](https://fluxml.ai/Flux.jl/stable/models/regularisation/#Flux.activations) outperforms [Zygote.Buffer](https://fluxml.ai/Zygote.jl/latest/utils/#Zygote.Buffer) significantly:

> 60.174614 seconds (79.67 M allocations: 16.958 GiB, 3.24% gc time)

vs

> 210.037932 seconds (691.79 M allocations: 49.948 GiB, 4.98% gc time)