# Flux + GPU memory problems

**URL:** <https://discourse.julialang.org/t/flux-gpu-memory-problems/79783>\
**Category:** Machine Learning\
**Tags:** flux\
**Created:** [April 21, 2022, 1:33pm UTC](https://discourse.julialang.org/t/flux-gpu-memory-problems/79783 "2022-04-21T13:33:18Z")\
**Posts on this page:** 3\
**Page:** 1

<div class="post-metadata">

**Author:** ![lgmendes](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/lgmendes/32/34577_2.png) [@lgmendes](https://discourse.julialang.org/u/lgmendes)\
**Post date:** [April 21, 2022, 1:33pm UTC](https://discourse.julialang.org/t/flux-gpu-memory-problems/79783/1 "2022-04-21T13:33:18Z")

</div>

I was trying to train some models available in the Metalhead (0.7) + Flux (0.13) using an NVIDIA GTX 1080 Ti (11 GB ) using some toy data (512 images size:224x224x3). I was running the code that I show at the end of the post using Metalhead.MobileNetv3(nclasses=2), Metalhead.ResNet34(pretrain=false,nclasses = 2) . However I was not able to run (memory problems) with Metalhead.ResNet50(pretrain=false,nclasses = 2) or Metalhead.ConvNeXt(:tiny, nclasses=2).

\*\*With other training platforms, I was able to train models like ResNet50. \*\*  
**Any problem with my training approach? Any advice?**

Thank you in advance!

```julia
device=gpu
#model=Metalhead.MobileNetv3(nclasses=2)|>device # OK!!!!!!!!!!!!!!!!!! 58% de memoria
#model =ResNet34(pretrain=false,nclasses = 2)|> device #OK!!!!!!!!!!!!!! 99% de memoria
model =ResNet50(pretrain=false,nclasses = 2)|> device #notOK

#model =Metalhead.ConvNeXt(:tiny, nclasses=2)|> device #notOK

loss(ŷ, y) = logitcrossentropy(ŷ, y)

ps = Flux.params(model);  

η = 3e-4 # learning rate
λ = 1e-6  
    #opt = ADAM(args.η)
opt = AdaBelief(η)
if λ > 0 # add weight decay, equivalent to L2 regularization
   opt = Optimiser(WeightDecay(λ), opt)
end
   
#GC.gc(true); CUDA.reclaim();  

for ind=1:10
@info ind
    for (x, y) in traindl

        y=onehotbatch(y, ["normal","abnormal"])
        x, y = x |> device, y |> device
        gs = Flux.gradient(ps) do
            ŷ = model(x)
            loss(ŷ, y)
        end

        Flux.Optimise.update!(opt, ps, gs)
    end

end

```

---

<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:** [April 21, 2022, 2:29pm UTC](https://discourse.julialang.org/t/flux-gpu-memory-problems/79783/2 "2022-04-21T14:29:51Z")

</div>

As a first step, you could wrap `traindl` in [`CUDA.CuIterator`](https://cuda.juliagpu.org/stable/usage/memory/#Batching-iterator). That should remove some memory pressure from the data loading.

---

<div class="post-metadata">

**Author:** ![lgmendes](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/lgmendes/32/34577_2.png) [@lgmendes](https://discourse.julialang.org/u/lgmendes)\
**Post date:** [April 26, 2022, 10:03pm UTC](https://discourse.julialang.org/t/flux-gpu-memory-problems/79783/3 "2022-04-26T22:03:54Z")

</div>

It seems that the example related to ImageNet that I have found in the ExplicitFluxLayers.jl [ImageNet Example]([Lux.jl/examples/ImageNet at main · avik-pal/Lux.jl · GitHub](https://github.com/avik-pal/ExplicitFluxLayers.jl/tree/main/examples/ImageNet) ExplicitFluxLayers.jl) addresses some problems related with GPU memory + Flux (that includes MPI; your suggestion is also used). The source code is also very well structured.