The problem: NetTrain’s custom BatchGenerationFunction runs synchronously on the same kernel that’s doing the training. If generating a batch is expensive, because it involves real work like loading and augmenting data or generated simulated data on the fly rather than indexing a pre-made array, every call blocks the GPU while the CPU does that work, then the CPU sits idle while the GPU trains. The two never overlap. If generation is slow enough it becomes the actual bottleneck, and a GPU that should run near-continuously ends up waiting on single-threaded CPU work between every batch. PyTorch solves this with DataLoader(num_workers=...), background worker processes that prefetch batches while the GPU is busy. As far as I can tell, Mathematica has no built-in equivalent.
The approach: build the missing piece with WL’s parallel kernels.
The approach: build the missing piece with WL’s parallel kernels.
◼
A pool of producer kernels independently generate examples (in real use: load and augment training data) and push finished examples into a shared queue, one slot per producer, so kernels never write to the same underlying list and step on each other.
◼
One dedicated trainer kernel runs NetTrain itself, with a lightweight BatchGenerationFunction that only pops ready-made examples off the queue. No expensive work happens on the kernel that’s actually training.
◼
The queue is bounded so producers can’t run ahead and blow up memory if they outpace consumption.
◼
Since there’s no built-in way to tell a running NetTrain “the data supply just died,” the batch function watches for that itself (all producers gone, or no data arriving within a timeout) and cleanly aborts the training call if it can no longer be fed, a safety net so the whole thing can’t hang indefinitely.
The test: a minimal synthetic network, with a random pause per example standing in for a slow, variable-cost augmentation step, deliberately CPU-bound rather than GPU-bound, to isolate the effect. Both runs given the same one-minute time budget:
◼
Baseline (generator running inline, as NetTrain normally works): 52 batches, 208 examples processed in the full minute, 3.4 examples/s.
◼
Parallel (producer kernels sized to the machine’s cores + 1 dedicated trainer kernel): 500 batches, 2000 examples, finished in 44 seconds, comfortably under budget, 46 examples/s.
◼
Same time allowance for both. The parallel version finished the entire training run; the baseline got through a tenth of it. That’s roughly a 1.x throughput improvement for identical per-example generation cost, purely from overlapping generation with training instead of serializing them.
A caveat worth being upfront about: because both a producer kernel and the trainer kernel touch the same underlying queue, one adding data, the other reading it, there’s a narrow window where a batch gets pulled at the exact moment the queue is being trimmed underneath it. In practice this can show up as a small fraction of examples occasionally being handed to training twice rather than being lost. With fast, uniform generation, the measured duplication rate was exactly zero in testing. With slower, more variable generation, a small rate reappeared, around 1%. Either way it doesn’t corrupt anything and it doesn’t stall training, it just means the network occasionally sees a duplicate augmented sample. For training data that’s already random and regenerated fresh every round, that’s a non-issue, nowhere near the concern that motivated this in the first place (a naive shared queue losing the majority of its updates under contention).
The duplication rate isn’t a fixed property of the design, it depends on how slow and variable the generation step is. The producer only checks whether to trim its queue once per loop pass, and that pass is gated behind however long generation takes. Slow generation lets the “already-delivered but not yet removed” backlog grow before it’s caught up, and a bigger catch-up trim is more likely to land in the brief window where the trainer is mid-read. Faster, more consistent generation keeps that backlog near zero and closes the race window on its own.
One more thing worth measuring: with generation delay removed entirely, total throughput was no faster than with a real delay per sample. With this many producer kernels, the coordination overhead itself (shared-variable synchronization, scanning across producer queues, kernel scheduling) has its own floor, around 0.2s per sample in this setup. Below that, the approach stops buying anything, since the bottleneck shifts from generation to orchestration. It’s only worth the added complexity when real augmentation cost is meaningfully above that floor, which for anything doing actual image loading, transformation, or on-the-fly augmentation is an easy bar to clear, but it’s a real limit worth knowing rather than assuming the speedup scales indefinitely as generation gets cheaper.
In[]:=
(*thenetwork*)net=NetChain[{LinearLayer[16],Ramp,LinearLayer[1]},"Input"->10];(*settingsandnetworkforTRAINER*)batchSize=4;roundLength=20;rounds=100;Echo["Starting parallel data generation"];(*Launchingthekernels*)Echo[DateString[],"Set up worker kernels"];CloseKernels[];LaunchKernels[];(*settingsforPRODUCER*)currentIndex=1;nProducers=Length[Kernels[]]-1;patchesPerAugment=3;maxQueue=3batchSize;maxRuns=2Ceiling[(roundLengthrounds+10)/(patchesPerAugmentnProducers)];activeProducers=nProducers;stallTimeout=4;Echo[{$KernelCount,nProducers},"Kernels {total, producers}:"];(*Getabatchfromthequeue*)batchFunction=Block[{q,bs,startTime,offset,attempt},bs=#BatchSize;startTime=AbsoluteTime[];(*loopoverallqueuestofindbatch*)Catch[While[True,Do[(*getthequeueandoffset*)q=queue[[currentIndex]];offset=takenSoFar[[currentIndex]]-trimmed[[currentIndex]];(*trytotakeabatchifqueueislongenough*)attempt=Quiet[Check[q[[offset+1;;offset+bs]],$Failed]];(*ifsuccesselectandremovethebatch*)If[attempt=!=$Failed,takenSoFar[[currentIndex]]+=bs;currentIndex=Mod[currentIndex,nProducers]+1;Throw[Merge[attempt,Identity]]];currentIndex=Mod[currentIndex,nProducers]+1;,{nProducers}];(*ifnobatchfoundinallqueuescheckwhyelsejustwaitabit*)If[activeProducers<=0,Echo["No more generators"];Abort[]];If[AbsoluteTime[]-startTime>stallTimeout,Echo["Could not get enough samples in 4s"];Abort[]];Pause[0.01]]]]&;(*functionthatactuallgeneratesthedatatobeaddedtothequeue*)(*placeholderforactualexpesivedatageneration/augmentation*)produceChunk[n_]:=Table[Pause[RandomReal[{0.05,.5}]];<|"Input"->RandomReal[1,10],"Output"->RandomReal[1,1]|>,{n}];(*updatethestatusformonotoringthetrainign*)(*canbemademuchmoreelaborateifneeded*)monitorFunction[snap_]:=status=Row[{"Round: ",Lookup[snap,"Round","?"]," / ","Batch loss: ",Lookup[snap,"BatchLoss","?"]}];(*monitoring*)trainingDone=False;status="waiting for first round...";queue=Table[{},nProducers];trueProduced=takenSoFar=trimmed=Table[0,nProducers];(*variablesneededduringevaluation,somesharedsomedistributed*)SetSharedVariable[trainingDone,status,activeProducers,trueProduced,queue,trimmed,takenSoFar];DistributeDefinitions[batchFunction,produceChunk,monitorFunction,maxQueue,maxRuns,batchSize,roundLength,rounds,patchesPerAugment];(*kernelmonitoring*)Echo[Dynamic[Length/@queue],"Queue lengths:"];Echo[Dynamic[activeProducers],"Active producers:"];Echo[Dynamic[Row[{!trainingDone," - ",status}]],"Training status:"];(*PRODUCERkernel:keepthequeuefeduntiltrainingsignalsdoneandremovesusedsamples*)producerJob=Table[With[{qi=i},ParallelSubmit[Block[{q,toTrim,wouldFit,c=0.0,localCount=0},(*checkiftrainingisbussy,seeifsamplesareneededdontgotofarahead*)While[!trainingDone&&localCount<maxRuns&&RandomChoice[{1-c,c}->{True,False}],q=queue[[qi]];toTrim=takenSoFar[[qi]]-trimmed[[qi]];wouldFit=Length[q]-toTrim<maxQueue;(*checkwhattodo,trimandappend,justtrim,ornothing*)Which[wouldFit,localCount++;queue[[qi]]=Join[Drop[q,toTrim],produceChunk[patchesPerAugment]];trimmed[[qi]]+=toTrim,toTrim>0,queue[[qi]]=Drop[q,toTrim];trimmed[[qi]]+=toTrim;Pause[0.01],True,Pause[0.01]]];(*ifproducersstopsforanyreasonsubtractit*)trueProduced[[qi]]=localCount;activeProducers--;]]],{i,nProducers}];(*TRAINERkernel:runNetTrain,signaldonewhenfinished*)trainerJob=With[{netI=net},ParallelSubmit[CheckAbort[Block[{result},result=NetTrain[netI,{batchFunction,"RoundLength"->roundLength},All,BatchSize->batchSize,MaxTrainingRounds->rounds,TargetDevice->"CPU",TimeGoal->Quantity[1,"Minutes"],TrainingProgressFunction->monitorFunction,TrainingProgressReporting->None];trainingDone=True;result],(trainingDone=True;$Aborted)]]];parallel=First[WaitAll[{trainerJob,producerJob}]];(*dosomevalidationaboutexampleslostorreused*)totalProduced=Total[trueProduced]*patchesPerAugment;totalRemaining=Total[Length/@queue];discrepancy=totalProduced-Total[takenSoFar]-totalRemaining;Echo[MatrixForm[{trueProduced,takenSoFar,trimmed}],"Runs vs True"];Echo[<|"Produced"->totalProduced,"Delivered"->Total[takenSoFar],"StillQueued"->totalRemaining,"Discrepancy"->discrepancy,"Effect"->Which[discrepancy>0,"examples missing",discrepancy<0,"examples duplicated",True,"exact"],"DiscrepancyPercent"->N[100discrepancy/totalProduced]|>,"Data loss audit:"];Echo[parallel,"The parallel evaluation"];Echo["Starting sequential data generation"];makeBatch=Merge[Table[Pause[RandomReal[{0.05,.5}]];<|"Input"->RandomReal[1,10],"Output"->RandomReal[1,1]|>,{#BatchSize}],Identity]&;sequential=NetTrain[net,{makeBatch,"RoundLength"->roundLength},All,BatchSize->batchSize,MaxTrainingRounds->rounds,TargetDevice->"CPU",TimeGoal->Quantity[1,"Minutes"]];Echo[sequential,"The sequential evaluation"];
CITE THIS NOTEBOOK
CITE THIS NOTEBOOK
Parallel batch generation overlapped with NetTrain via parallel kernels
by Martijn Froeling
Wolfram Community, STAFF PICKS, August 29, 2026
https://community.wolfram.com/groups/-/m/t/3789292
by Martijn Froeling
Wolfram Community, STAFF PICKS, August 29, 2026
https://community.wolfram.com/groups/-/m/t/3789292