# Avoid allocation of a Flux model on the CPU

**URL:** <https://discourse.julialang.org/t/avoid-allocation-of-a-flux-model-on-the-cpu/121785>\
**Category:** Machine Learning\
**Tags:** flux\
**Created:** [October 26, 2024, 1:14pm UTC](https://discourse.julialang.org/t/avoid-allocation-of-a-flux-model-on-the-cpu/121785 "2024-10-26T13:14:49Z")\
**Posts on this page:** 4\
**Page:** 1

<div class="post-metadata">

**Author:** ![mesonepigreco](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mesonepigreco/32/46262_2.png) [@mesonepigreco](https://discourse.julialang.org/u/mesonepigreco)\
**Post date:** [October 26, 2024, 1:14pm UTC](https://discourse.julialang.org/t/avoid-allocation-of-a-flux-model-on-the-cpu/121785/1 "2024-10-26T13:14:49Z")

</div>

Hi everyone,  
I have to execute a flux model inside a Monte Carlo simulation. I am currently working on the CPU; I am facing the problem of executing a `model(configuration)`for each step of the Monte Carlo, which allocates memory. Since this runs on a nonparallelizable loop (each iteration depends on the result of the previous one, so using batches is not a solution), I got tons of allocations that make the GC hit continuously.  
By inspecting the Profiling of the allocation, I get the following fireplot

 ![image](https://global.discourse-cdn.com/julialang/original/3X/f/2/f2b8f634edf9e22b68ec2dd6074601505c381f44.png)

It seems that most of the allocation here happens in the Convolutional layers of the neural network.

The model is defined as

```julia
model = Chain(x -> 2x .- 1,
          x -> reshape(x, (8, 8, 1,1)),
		  Conv((3,3), 1=>4, tanh; pad=1),
		  Conv((3,3), 4=>8, tanh; pad=1),
		  x -> reshape(x, :),
		  Dense(8*64, n_observables, tanh))

```

And the function `monte_carlo_run!` calls the model continuously inside a `for` loop.

How can I avoid the allocations inside the convolutioinnal layers?

---

<div class="post-metadata">

**Author:** ![Tomas\_Pevny](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/tomas_pevny/32/25466_2.png) [@Tomas\_Pevny](https://discourse.julialang.org/u/Tomas_Pevny)\
**Post date:** [October 26, 2024, 3:08pm UTC](https://discourse.julialang.org/t/avoid-allocation-of-a-flux-model-on-the-cpu/121785/2 "2024-10-26T15:08:17Z")

</div>

If it is only for inference, you can try this [GitHub - ericphanson/AllocArrays.jl: Arrays that use a dynamically-scoped allocator](https://github.com/ericphanson/AllocArrays.jl). If it works, report the results.

---

<div class="post-metadata">

**Author:** ![CarloLucibello](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/carlolucibello/32/3278_2.png) [@CarloLucibello](https://discourse.julialang.org/u/CarloLucibello)\
**Post date:** [October 27, 2024, 5:04am UTC](https://discourse.julialang.org/t/avoid-allocation-of-a-flux-model-on-the-cpu/121785/3 "2024-10-27T05:04:43Z")

</div>

If you don’t have to train the network, you can define your own custom Conv layer that calls the non-allocating `NNlib.conv!(y, x, w, cdims)`. See here:

> <https://github.com/FluxML/NNlib.jl/blob/master/src/conv.jl>

---

<div class="post-metadata">

**Author:** ![mesonepigreco](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mesonepigreco/32/46262_2.png) [@mesonepigreco](https://discourse.julialang.org/u/mesonepigreco)\
**Post date:** [October 27, 2024, 8:12am UTC](https://discourse.julialang.org/t/avoid-allocation-of-a-flux-model-on-the-cpu/121785/4 "2024-10-27T08:12:19Z")

</div>

Wow, AllocArrays seems to effectively reduce by far the memory allocated by Flux with almost zero effort. Some memory is still allocated but reduced by a factor of 50. I will implement it in the actual program and see how it goes.

```julia
using Bumper, Flux, AllocArrays

x = zeros(Float32, 32, 32, 1, 1)
x_new = AllocArray(x)

model = Chain(Conv((3,3), 1=>4, relu; pad=1), 
              x ->reshape(x, :), 
              Dense(4*32*32=>4))

function simple_run(model, data, iterations)
       result = model(data)
       tmp_input = similar(data)
       
       for i in 2:iterations
           tmp_input .= data
           tmp_input .+= i
           result .+= model(tmp_input)
       end
       return result
end

function bumper_run(model, data, iterations)
       b = UncheckedBumperAllocator(2^20)
       result = model(data)
       tmp_input = similar(data)
       with_allocator(b) do
       for i in 2:iterations
           tmp_input .= data
           tmp_input .+= i
           result .+= model(tmp_input)
           reset!(b)
       end
       end
       return result
end

@time simple_run(model, x, 1)
# 0.000129 seconds (58 allocations: 76.188 KiB)

@time simple_run(model, x, 10000)
# 0.501996 seconds (560.00 k allocations: 703.434 MiB, 1.98% gc time)

@time bumper_run(model, x, 1)
# 0.000151 seconds (79 allocations: 1.075 MiB)

@time bumper_run(model, x, 10000)
# 0.554891 seconds (580.02 k allocations: 37.540 MiB, 1.21% gc time)

```