Training FLUX models with larger datasets


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:

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)
            Flux.Optimise.update!(opt, ps, gs)

The inputs of the DataLoader are

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

I found a package of a DataLoader implementation (GitHub - lorenzoh/DataLoaders.jl: A parallel iterator for large machine learning datasets that don't fit into memory inspired by PyTorch's `DataLoader` class.) that works like PyTorch’s DataLoader which seems to solve the main problem!

1 Like

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


From the DataLoaders documentation (DataLoaders.jl/ at master · lorenzoh/DataLoaders.jl · GitHub) 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

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

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:

import DataLoaders.LearnBase: getobs, nobs
using Images

struct ImageDataset
ImageDataset(folder::String) = ImageDataset(readdir(folder))

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

data = ImageDataset("path/to/my/images")
for images in DataLoader(data, 16)
    # Do something

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.