%% Use minibatchqueue in a Custom Training Loop
% Train a network using minibatchqueue to manage the processing of mini-batches.
%% Load Training Data
% Load the digits training data and store the data in a datastore. Create a datastore for the images and one for the labels using arrayDatastore. Then, combine the datastores to produce a single datastore to use with minibatchqueue.

[XTrain,YTrain] = digitTrain4DArrayData;
dsX = arrayDatastore(XTrain,'IterationDimension',4);
dsY = arrayDatastore(YTrain);

dsTrain = combine(dsX,dsY);

% Determine the number of unique classes in the label data.

classes = categories(YTrain);
numClasses = numel(classes);

%
%% Define Network
% Define the network and specify the average image value using the 'Mean' option in the image input layer.

layers = [
    imageInputLayer([28 28 1], 'Name','input','Mean',mean(XTrain,4))
    convolution2dLayer(5,20,'Name','conv1')
    reluLayer('Name', 'relu1')
    convolution2dLayer(3,20,'Padding',1,'Name','conv2')
    reluLayer('Name','relu2')
    convolution2dLayer(3,20,'Padding',1,'Name','conv3')
    reluLayer('Name','relu3')
    fullyConnectedLayer(numClasses,'Name','fc')
    softmaxLayer('Name','softmax')];
lgraph = layerGraph(layers);

% Create a dlnetwork object from the layer graph.

dlnet = dlnetwork(lgraph);

%
%% Define Model Gradients Function
% Create the helper function modelGradients, listed at the end of the example. The function takes a dlnetwork object dlnet and a mini-batch of input data dlX with corresponding labels Y, and returns the loss and the gradients of the loss with respect to the learnable parameters in dlnet.
%
%% Specify Training Options
% Specify the options to use during training.

numEpochs = 10;
miniBatchSize = 128;

% Visualize the training progress in a plot.

plots = "training-progress";

%
%% Create the minibatchqueue
% Use minibatchqueue to process and manage the mini-batches of images. For each mini-batch:
% Discard partial mini-batches.
% Use the custom mini-batch preprocessing function preprocessMiniBatch (defined at the end of this example) to one-hot encode the class labels.
% Format the image data with the dimension labels 'SSCB' (spatial, spatial, channel, batch). By default, the minibatchqueue object converts the data to dlarray objects with underlying type single. Do not add a format to the class labels.
% Train on a GPU if one is available. By default, the minibatchqueue object converts each output to a gpuArray if a GPU is available. Using a GPU requires Parallel Computing Toolbox™ and a CUDA® enabled NVIDIA® GPU with compute capability 3.0 or higher.

mbq = minibatchqueue(dsTrain,...
    'MiniBatchSize',miniBatchSize,...
    'PartialMiniBatch','discard',...
    'MiniBatchFcn',@preprocessMiniBatch,...    
    'MiniBatchFormat',{'SSCB',''});

%
%% Train Network
% Train the model using a custom training loop. For each epoch, shuffle the data and loop over mini-batches while data is still available in the minibatchqueue. Update the network parameters using the adamupdate function. At the end of each epoch, display the training progress.
% Initialize the training progress plot.

if plots == "training-progress"
    figure
    lineLossTrain = animatedline('Color',[0.85 0.325 0.098]);
    ylim([0 inf])
    xlabel("Iteration")
    ylabel("Loss")
    grid on
end

% Initialize the average gradients and squared average gradients.

averageGrad = [];
averageSqGrad = [];

% Train the network.

iteration = 0;
start = tic;

for epoch = 1:numEpochs
    % Shuffle data.
    shuffle (mbq);
        
    while hasdata(mbq)
        iteration = iteration + 1;
        
        % Read mini-batch of data
        [dlX,Y] = next(mbq);
              
        % Evaluate the model gradients and loss using dlfeval and the
        % modelGradients helper function.
        [grad,loss] = dlfeval(@modelGradients,dlnet,dlX,Y);
        
        % Update the network parameters using the Adam optimizer.
        [dlnet,averageGrad,averageSqGrad] = adamupdate(dlnet,grad,averageGrad,averageSqGrad,iteration);
        
        % Display the training progress.
        if plots == "training-progress"
            D = duration(0,0,toc(start),'Format','hh:mm:ss');
            addpoints(lineLossTrain,iteration,double(gather(extractdata(loss))))
            title("Epoch: " + epoch + ", Elapsed: " + string(D))
            drawnow
        end
    end
end

%
%% Model Gradients Function
% The modelGradients helper function takes a dlnetwork object dlnet and  a mini-batch of input data dlX with corresponding labels Y, and returns the loss and the gradients of the loss with respect to the learnable parameters in dlnet. To compute the gradients automatically, use the dlgradient function.

function [gradients,loss] = modelGradients(dlnet,dlX,Y)
    dlYPred = forward(dlnet,dlX);    
    loss = crossentropy(dlYPred,Y);    
    gradients = dlgradient(loss,dlnet.Learnables);
    
end

%% Mini-Batch Preprocessing Function
% The preprocessMiniBatch function preprocesses the data using the following steps:
% Extract the image data from the incoming cell array and concatenate into a numeric array. Concatenating the image data over the fourth dimension adds a third dimension to each image, to be used as a singleton channel dimension.
% Extract the label data from the incoming cell array and concatenate along the second dimension into a categorical array.
% One-hot encode the categorical labels into numeric arrays. Encoding into the first dimension produces an encoded array that matches the shape of the network output.


function [X,Y] = preprocessMiniBatch(XCell,YCell)
    % Extract image data from cell and concatenate over 4th dimension to adds a
    % singleton dimension 3 for channel dimension
    X = cat(4,XCell{:});

    % Extract label data from cell and concatenate
    Y = cat(2,YCell{:});
    
    % One-hot encode labels
    Y = onehotencode(Y,1);

end

% Copyright 2020 The MathWorks, Inc.
%