How to NetTrain only N batches?

Question

I'd like to train a neural network for only a specified number of batches, that is for only part of an epoch (or TrainingRound). I've tried using the NetTrain[net,{f,"RoundLength"->n}] form, but this only produces approximately n examples.

n = 10000;
trainingData = RandomReal[1, {n, 4}] -> RandomReal[1, {n, 4}];
net = NetChain[{8, Tanh, 2048, Tanh, 2048, Tanh, 4}];
trained = NetTrain[net, trainingData, BatchSize-> 24, 
             MaxTrainingRounds -> 0.1 (* I'd like 'TrainingBatches -> 6' *)]

Answer 1

The function NetTrain will build a NetTrainer first.

Then NeuralNetworks`TrainerUpdate[trainer, DataBatch] will update the net.

Let's see the options of NeuralNetworks`ToNetTrainer

Options[NeuralNetworks`ToNetTrainer]

{
    "BatchSize" -> Automatic, "MaxBatchSize" -> None,
    "TotalBatches" -> 4096, "UpdatesPerBatch" -> 1,
    "Context" -> 1, "DataType" -> 0, "GradientScale" -> 1,
    "Optimizer" -> "ADAM", "LearningRates" -> Automatic
}

"TotalBatches" is an optional item, so I think your needs are in development.

See if the next version will add this option.

---

Well, in fact, we do not need to know these.

Use "StopTraining", and if you want to store each of the 6 batches, using "Checkpointing".

n = 1000;
trainingData = RandomReal[1, {n, 4}] -> RandomReal[1, {n, 4}];
net = NetChain[{8, Tanh, 2048, Tanh, 2048, Tanh, 4}];
trained = NetTrain[net, trainingData, All, TargetDevice -> "GPU",
    BatchSize -> 24, MaxTrainingRounds -> 10000,
    TrainingProgressFunction -> {"StopTraining"&, "Interval" -> Quantity[60, "Batches"]},
    TrainingProgressCheckpointing -> {"Directory", $HomeDirectory, "Interval" -> Quantity[6, "Batches"]}
]

If you want to store the specific {2, 3, 5, 7} batch's training results.

Using NetTrainResultsObject, and extract the final result and repeat the first step.