Unexpected high memory consumption using out-of-core NetTrain

Question

I'm working on training a neural network on an image dataset. There are 14k images and each image contains 3x150x150 pixels. I have built a generator function following the approach in Training on Large Datasets for Out-of-Core training. According to the reference

The first approach is for users to write a generator function f that, when evaluated, can load a single batch of data from an external source such as a disk or database. NetTrain[net,f,…] calls f at each training batch iteration, thus only keeping a single batch of training data in memory.

I use a generator function in that documentation:

genTrain = Function[RandomSample[trainingData, #BatchSize]]

The output of it is as expected:

However, the memory consumption of NetTrain keeps growing every round. I use a small RoundLength here in case of memory exhaust:

trainedNet = NetTrain[lenet, {genTrain, "RoundLength" -> 280}, 
TargetDevice -> "GPU", MaxTrainingRounds -> 50, BatchSize -> 64, 
WorkingPrecision -> "Mixed", PerformanceGoal -> "TrainingMemory"]

Here is the memory usage from the task manager:

It increases till the end.

I've tried $HistoryLength=0, Module[{x=Function[RandomSample[trainingData, #BatchSize]]},x] but nothing changed.

My questions are:

1. Am I getting the right idea about Out-of-Core training? It should load a single batch of data and free it after use, isn't it? Or is there a bug?

2. What can I do to decrease memory consumption, for example, free the memory after every round or batch?

Answer 1

As Charmbracelet mentioned in the comments, this bug was fixed in 13.0.1 but just in case it's useful here is a way to work around it based on the generator approach you mentioned.

LoadTrainingPair[idx_] := Module[{out},
  out = <|"Input" -> 
     First@Image`ImportExportDump`ImageReadPNG[frogImages[[idx]]], 
    "Output" -> 
     First@Image`ImportExportDump`ImageReadPNG[gtImages[[idx]]]|>;
  Image`ImportExportDump`DeleteCachePNG[];
  out
  ]

You can see the full story + example notebook with some helpful feedback from Wolfram folks on the community forum (in my case I was working on per-pixel segmentation so the memory leakage was twice as bad).