%% Image Classification Using Custom Training Loop
% Use this training function to define the training data, network architecture, 
% training options, and custom training loop used by the experiment. Experiment 
% Manager saves the output of this function, so you can export it to the MATLAB 
% workspace when the training is complete. For more information on using this 
% template, see <matlab:helpview('deeplearning','ExperimentManagerAppCustomTrainingExample') 
% Image Classification Using Custom Training Loop>. For more information on custom 
% experiments, see <matlab:helpview('deeplearning','exp-mgr-create-custom-experiment') 
% Configure Custom Training Experiment>.
% 
% Modify this function as a template to set up your own experiment.
%% Input
%% 
% * |params| is a structure with fields from the Experiment Manager hyperparameter 
% table.
% * |monitor| is an <matlab:doc('experiments.Monitor') experiments.Monitor> 
% object that you can use to track the progress of the training, update information 
% fields in the results table, record values of the metrics used by the training, 
% and produce training plots.
%% Output
%% 
% * |output| is a structure that contains the trained network and the values 
% of the training loss.
function output = {functionName}(params,monitor)
%% Load Training Data
% Load the training data for the experiment as 4-D arrays. The training data 
% consists of 5000 images from the Digits data set. Each image shows a digit from 
% 0 to 9, rotated by a certain angle.
[XTrain,TTrain] = digitTrain4DArrayData;
classes = categories(TTrain);
numClasses = numel(classes);
%% 
% Replace the data with your own training data. For more information, see <matlab:helpview('deeplearning','data-sets-for-deep-learning') 
% Data Sets for Deep Learning>.
%% Define Network Architecture
% Define the convolutional neural network architecture in the layer graph |layers|. 
% To train the network with a custom training loop, convert the layer graph to 
% a <matlab:doc('dlnetwork') |dlnetwork|> object.
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')];
lgraph = layerGraph(layers);
dlnet = dlnetwork(lgraph);
%% 
% Replace the layer graph with your own layer graph to train. For more information, 
% see <matlab:helpview('deeplearning','list_of_layers') List of Deep Learning 
% Layers>.
%% Specify Training Options
% Read the hyperparameters specified in Experiment Manager by accessing the 
% fields of the input |params|. For example, access the initial learning rate 
% |myInitialLearnRate| defined in the hyperparameter table with the argument value 
% |params.myInitialLearnRate|.
velocity = [];
numEpochs = 20;
miniBatchSize = 128;
numObservations = numel(TTrain);
numIterationsPerEpoch = floor(numObservations./miniBatchSize);
initialLearnRate = params.myInitialLearnRate;
momentum = 0.9;
decay = 0.01;
executionEnvironment = "auto";
%% Train Model
% Train the model using a custom training loop. Loop over the training data 
% and update the network parameters at each iteration. Convert the data to <matlab:doc('dlarray') 
% |dlarray|> objects with underlying type |single| and specify the dimension labels 
% |'SSCB'| (spatial, spatial, channel, batch) for the image data. For GPU training, 
% convert the data to <matlab:doc('gpuArray') |gpuArray|> objects. Evaluate the 
% model gradients and loss using <matlab:doc('dlfeval') |dlfeval|> and the helper 
% function |modelGradients|. Determine the learning rate using a time-based decay 
% learning rate schedule. Update the network parameters using the SGDM optimizer. 
% 
% Use the input argument |monitor| to track the progress of the training and 
% record the value of the training loss. After each iteration of the custom training 
% loop, save the trained network and update training progress. Exit the custom 
% training loop if you click *Stop* in the Experiment Manager toolstrip. 
iteration = 0;
monitor.Metrics = "TrainingLoss";
monitor.XLabel = "Iteration";
output.trainingInfo.loss = [];
for epoch = 1:numEpochs
    idx = randperm(numel(TTrain));
    XTrain = XTrain(:,:,:,idx);
    TTrain = TTrain(idx);
    for i = 1:numIterationsPerEpoch
        iteration = iteration + 1;
        idx = (i-1)*miniBatchSize+1:i*miniBatchSize;
        X = XTrain(:,:,:,idx);
        
        Y = zeros(numClasses, miniBatchSize, 'single');
        for c = 1:numClasses
            Y(c,TTrain(idx)==classes(c)) = 1;
        end
        
        dlX = dlarray(single(X),'SSCB');
        
        if (executionEnvironment == "auto" && canUseGPU) || executionEnvironment == "gpu"
            dlX = gpuArray(dlX);
        end
        
        [gradients,loss] = dlfeval(@modelGradients,dlnet,dlX,Y);
        
        learnRate = initialLearnRate/(1 + decay*iteration);
        
        [dlnet.Learnables, velocity] = sgdmupdate(dlnet.Learnables, gradients, velocity, learnRate, momentum);
        
        curLoss = double(gather(extractdata(loss)));
        recordMetrics(monitor,iteration,'TrainingLoss',curLoss);
        
        output.trainingInfo.loss = [output.trainingInfo.loss; iteration, curLoss];
        
    end    
    output.trainedNet = dlnet;
    monitor.Progress = (epoch*100.0)/numEpochs;
    
    if monitor.Stop
        return;
    end
end
%% 
% End of training function.
end
%% 
% Helper function to evaluate the model gradients and loss.
function [gradients,loss] = modelGradients(dlnet,dlX,Y)
dlYPred = forward(dlnet,dlX);
dlYPred = softmax(dlYPred);
loss = crossentropy(dlYPred,Y);
gradients = dlgradient(loss,dlnet.Learnables);
end
%% 
% Copyright 2020-2021 The MathWorks, Inc.