# Training FLUX models with larger datasets

**URL:** <https://discourse.julialang.org/t/training-flux-models-with-larger-datasets/79057>\
**Category:** Machine Learning\
**Tags:** cuda, flux\
**Created:** [April 5, 2022, 3:29pm UTC](https://discourse.julialang.org/t/training-flux-models-with-larger-datasets/79057 "2022-04-05T15:29:11Z")\
**Posts on this page:** 5\
**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 5, 2022, 3:29pm UTC](https://discourse.julialang.org/t/training-flux-models-with-larger-datasets/79057/1 "2022-04-05T15:29:11Z")

</div>

Hi,

I’m trying to train a FLUX model (ResNet50) with a medium-size dataset (~ 30K images (512X512X3)) on a computer with 128GB of RAM and Three NVIDIA GeForce GTX 1080 Ti.

I’m loading the data with Dataloader using a BatchSize=32. When I try to process all → it runs out of memory. Even with a subset of the data, easily goes out of memory or the memory does not fit in the GPU  
Any advice on how to handle the problem?

Thanks in advance!

The training Loop is something like:

```julia
for epoch in 1:args.epochs
        @showprogress for (x, y) in train_loader
            x, y = x |> device, y |> device
            gs = Flux.gradient(ps) do
                    ŷ = model(x)
                    loss(ŷ, y)
                end
            Flux.Optimise.update!(opt, ps, gs)
 end

```

The inputs of the DataLoader are

```julia
octDfTraningImagesResize=map(x->imresize(RGB.(x),(imageResizeDim,imageResizeDim))|> channelview.|> Float32,octDfTraningImages)
 octDfTraningImagesResizeVect=map(x->permutedims(x,(3, 2, 1)),octDfTraningImagesResize)

```

---

<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 5, 2022, 10:33pm UTC](https://discourse.julialang.org/t/training-flux-models-with-larger-datasets/79057/2 "2022-04-05T22:33:40Z")

</div>

I found a package of a DataLoader implementation ([https://github.com/lorenzoh/DataLoaders.jl](https://github.com/lorenzoh/DataLoaders.jl)) that works like PyTorch’s `DataLoader` which seems to solve the main problem!

---

<div class="post-metadata">

**Author:** ![Rasmus\_Hoier](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/rasmus_hoier/32/24036_2.png) [@Rasmus\_Hoier](https://discourse.julialang.org/u/Rasmus_Hoier)\
**Post date:** [April 6, 2022, 7:07am UTC](https://discourse.julialang.org/t/training-flux-models-with-larger-datasets/79057/3 "2022-04-06T07:07:14Z")

</div>

MLUtils.jl might be interesting for you as well. I think the DataLoader in MLUtils is derived from the one in DataLoaders.jl.  
[MLUtils documentation](https://juliaml.github.io/MLUtils.jl/dev/)

---

<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 6, 2022, 12:07pm UTC](https://discourse.julialang.org/t/training-flux-models-with-larger-datasets/79057/4 "2022-04-06T12:07:32Z")

</div>

Thanks!

From the DataLoaders documentation ([https://github.com/lorenzoh/DataLoaders.jl/blob/master/docs/datacontainers.md](https://github.com/lorenzoh/DataLoaders.jl/blob/master/docs/datacontainers.md)) it is easy to create a dataloader object with the links and the methods to (lazy) read the images and the labels. However, how can I shuffle the data? the shuffleobs function does not work with “customized” dataTypes `DataLoaders.DataLoader(shuffleobs(data), 16) not work`

```julia
shuffleobs(data) = shuffleobs(Random.GLOBAL_RNG, data)
function shuffleobs(rng::AbstractRNG, data)
    obsview(data, randperm(rng, numobs(data)))
end

```

**Any suggestion?**

Possible practical solution may be , for each epoch, create/update the train\_loader with a shuffle version of the links to the images. Since we are working only with links this operation should be fast. A more elegant solution may be possible and/or already implemented.

The example code is:

```julia
import DataLoaders.LearnBase: getobs, nobs
using Images

struct ImageDataset
    files::Vector{String}
end
ImageDataset(folder::String) = ImageDataset(readdir(folder))

nobs(data::ImageDataset) = length(data.files)
getobs(data::ImageDataset, i::Int) = Images.load(data.files[i])

```

```julia

data = ImageDataset("path/to/my/images")
for images in DataLoader(data, 16)
    # Do something
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 7, 2022, 1:45am UTC](https://discourse.julialang.org/t/training-flux-models-with-larger-datasets/79057/5 "2022-04-07T01:45:44Z")

</div>

In cases like this, the custom dataset type is absolutely the way to go. If you’re looking for additional inspiration there, check out what FastAI.jl does for its image datasets.