Customize Output During Deep Learning Network Training

This example shows how to define an output function that runs at each iteration during training of deep learning neural networks. If you specify output functions by using the 'OutputFcn' name-value pair argument of trainingOptions, then trainNetwork calls these functions once before the start of training, after each training iteration, and once after training has finished. Each time the output functions are called, trainNetwork passes a structure containing information such as the current iteration number, loss, and accuracy. You can use output functions to display or plot progress information, or to stop training. To stop training early, make your output function return true. If any output function returns true, then training finishes and trainNetwork returns the latest network.

To stop training when the loss on the validation set stops decreasing, simply specify validation data and a validation patience using the 'ValidationData' and the 'ValidationPatience' name-value pair arguments of trainingOptions, respectively. The validation patience is the number of times that the loss on the validation set can be larger than or equal to the previously smallest loss before network training stops. You can add additional stopping criteria using output functions. This example shows how to create an output function that stops training when the classification accuracy on the validation data stops improving. The output function is defined at the end of the script.

Load the training data, which contains 5000 images of digits. Set aside 1000 of the images for network validation.

[XTrain,YTrain] = digitTrain4DArrayData;

idx = randperm(size(XTrain,4),1000);
XValidation = XTrain(:,:,:,idx);
XTrain(:,:,:,idx) = [];
YValidation = YTrain(idx);
YTrain(idx) = [];

Construct a network to classify the digit image data.

layers = [
    imageInputLayer([28 28 1])
    
    convolution2dLayer(3,8,'Padding','same')
    batchNormalizationLayer
    reluLayer   
    
    maxPooling2dLayer(2,'Stride',2)
    
    convolution2dLayer(3,16,'Padding','same')
    batchNormalizationLayer
    reluLayer   
    
    maxPooling2dLayer(2,'Stride',2)
    
    convolution2dLayer(3,32,'Padding','same')
    batchNormalizationLayer
    reluLayer   
    
    fullyConnectedLayer(10)
    softmaxLayer
    classificationLayer];

Specify options for network training. To validate the network at regular intervals during training, specify validation data. Choose the 'ValidationFrequency' value so that the network is validated once per epoch.

To stop training when the classification accuracy on the validation set stops improving, specify stopIfAccuracyNotImproving as an output function. The second input argument of stopIfAccuracyNotImproving is the number of times that the accuracy on the validation set can be smaller than or equal to the previously highest accuracy before network training stops. Choose any large value for the maximum number of epochs to train. Training should not reach the final epoch because training stops automatically.

miniBatchSize = 128;
validationFrequency = floor(numel(YTrain)/miniBatchSize);
options = trainingOptions('sgdm', ...
    'InitialLearnRate',0.01, ...
    'MaxEpochs',100, ...
    'MiniBatchSize',miniBatchSize, ...
    'VerboseFrequency',validationFrequency, ...
    'ValidationData',{XValidation,YValidation}, ...
    'ValidationFrequency',validationFrequency, ...
    'Plots','training-progress', ...
    'OutputFcn',@(info)stopIfAccuracyNotImproving(info,3));

Train the network. Training stops when the validation accuracy stops increasing.

net = trainNetwork(XTrain,YTrain,layers,options);

Training on single CPU.
Initializing input data normalization.
|======================================================================================================================|
|  Epoch  |  Iteration  |  Time Elapsed  |  Mini-batch  |  Validation  |  Mini-batch  |  Validation  |  Base Learning  |
|         |             |   (hh:mm:ss)   |   Accuracy   |   Accuracy   |     Loss     |     Loss     |      Rate       |
|======================================================================================================================|
|       1 |           1 |       00:00:04 |        7.81% |       12.70% |       2.7155 |       2.5169 |          0.0100 |
|       1 |          31 |       00:00:11 |       71.09% |       75.30% |       0.8804 |       0.8147 |          0.0100 |
|       2 |          62 |       00:00:17 |       88.28% |       88.20% |       0.3876 |       0.4419 |          0.0100 |
|       3 |          93 |       00:00:23 |       94.53% |       94.10% |       0.2181 |       0.2536 |          0.0100 |
|       4 |         124 |       00:00:31 |       96.09% |       96.70% |       0.1450 |       0.1746 |          0.0100 |
|       5 |         155 |       00:00:37 |       99.22% |       97.50% |       0.1001 |       0.1299 |          0.0100 |
|       6 |         186 |       00:00:44 |       99.22% |       97.90% |       0.0773 |       0.1118 |          0.0100 |
|       7 |         217 |       00:00:49 |      100.00% |       98.10% |       0.0554 |       0.0932 |          0.0100 |
|       8 |         248 |       00:00:54 |      100.00% |       98.00% |       0.0436 |       0.0861 |          0.0100 |
|       9 |         279 |       00:01:04 |      100.00% |       98.00% |       0.0340 |       0.0772 |          0.0100 |
|      10 |         310 |       00:01:09 |      100.00% |       98.50% |       0.0273 |       0.0671 |          0.0100 |
|      11 |         341 |       00:01:15 |      100.00% |       98.50% |       0.0236 |       0.0609 |          0.0100 |
|      12 |         372 |       00:01:22 |      100.00% |       98.90% |       0.0211 |       0.0556 |          0.0100 |
|      13 |         403 |       00:01:28 |      100.00% |       98.90% |       0.0186 |       0.0521 |          0.0100 |
|      14 |         434 |       00:01:33 |      100.00% |       98.90% |       0.0162 |       0.0497 |          0.0100 |
|      15 |         465 |       00:01:39 |      100.00% |       98.80% |       0.0143 |       0.0474 |          0.0100 |
|======================================================================================================================|
Training finished: Stopped by OutputFcn.

Define Output Function

Define the output function stopIfAccuracyNotImproving(info,N), which stops network training if the best classification accuracy on the validation data does not improve for N network validations in a row. This criterion is similar to the built-in stopping criterion using the validation loss, except that it applies to the classification accuracy instead of the loss.

function stop = stopIfAccuracyNotImproving(info,N)

stop = false;

% Keep track of the best validation accuracy and the number of validations for which
% there has not been an improvement of the accuracy.
persistent bestValAccuracy
persistent valLag

% Clear the variables when training starts.
if info.State == "start"
    bestValAccuracy = 0;
    valLag = 0;
    
elseif ~isempty(info.ValidationLoss)
    
    % Compare the current validation accuracy to the best accuracy so far,
    % and either set the best accuracy to the current accuracy, or increase
    % the number of validations for which there has not been an improvement.
    if info.ValidationAccuracy > bestValAccuracy
        valLag = 0;
        bestValAccuracy = info.ValidationAccuracy;
    else
        valLag = valLag + 1;
    end
    
    % If the validation lag is at least N, that is, the validation accuracy
    % has not improved for at least N validations, then return true and
    % stop training.
    if valLag >= N
        stop = true;
    end
    
end

end