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)
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)
MLUtils.jl might be interesting for you as well. I think the DataLoader in MLUtils is derived from the one in DataLoaders.jl.
From the DataLoaders documentation (DataLoaders.jl/datacontainers.md 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)))
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
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.