function output = {functionName}(params,monitor)
%%  Audio 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 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. This template uses two parameters: The type of auditory
% spectrogram you extract from the audio signal and the initial learning rate.
% * |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.
%% Load Training Data
% Download the Free Spoken Digit Data Set (FSDD). FSDD consists of 2000 recordings 
% of four speakers saying the numbers 0 through 9 in English.
downloadFolder = matlab.internal.examples.downloadSupportFile("audio","FSDD.zip");
dataFolder = tempdir;
unzip(downloadFolder,dataFolder)
dataset = fullfile(dataFolder,"FSDD");
%% 
% Create an <matlab:doc('audioDatastore') |audioDatastore|> that points to the dataset.
ads = audioDatastore(dataset,IncludeSubfolders=true);
%% 
% Decode the file names to set the labels (spoken digits) on the datastore.
ads.Labels = filenames2labels(ads.Files,ExtractBefore="_");
%% 
% Call <matlab:doc('audioDatastore.splitEachLabel') |splitEachLabel|>  to
% split the FSDD into training and test sets. You use the training set to
% train the model and the test set to validate the trained model. Allocate
% 80% of the data to the training set and retain 20% for the test set.
[adsTrain,adsTest] = splitEachLabel(ads,0.8,"randomized");
%% 
% To limit the runtime of the experiment, reduce the size of the training
% set. To train on the entire training set, remove this line.
adsTrain = splitEachLabel(adsTrain,0.75,"randomized");
%% Define Feature Extraction
% The deep learning network accepts auditory spectrograms.
% 
% Define parameters used to extract auditory spectrograms. Use 220 ms
% windows with 10 ms hops between windows. Use a 2048-point DFT and 40
% frequency bands. The sample rate of the dataset is 8 kHz. If you use
% another dataset with a different sample rate, change the value of |fs|
% accordingly.
frameDuration = 0.22;
hopDuration = 0.01;
fftLength = 2048;
numBands = 40;
fs = 8000;
frameLength = round(frameDuration*fs);hopLength = round(hopDuration*fs);
segmentLength = 8192;
%% 
% Create an <matlab:doc('audioFeatureExtractor') |audioFeatureExtractor|> 
% object to compute auditory spectrograms from input audio signals.
afe = audioFeatureExtractor(SampleRate=fs, ...
    Window=hamming(frameLength,"periodic"),OverlapLength=frameLength - hopLength, ...
    FFTLength=fftLength);
%% 
% Access the type of auditory spectrogram defined in the hyperparameter
% table with the argument value |params.spectrogramType|. Set the computed
% audio feature to the desired spectrum type.
switch params.spectrogramType
    case "mel"
        spect = "melSpectrum";
    case "bark"
        spect = "barkSpectrum";
    case "erb"
        spect = "erbSpectrum";
end
afe.(spect) = true;
%% 
% Set the parameters for the auditory spectrogram.
setExtractorParameters(afe,spect,NumBands=numBands,FrequencyRange=[50 fs/2]);
%% 
% Create a transformed datastore that computes mel-frequency spectrograms from 
% audio data. The supporting function, |getSpectrogram|, standardizes the 
% recording length and normalizes the amplitude of the audio input. |getSpectrogram| 
% uses the |audioFeatureExtractor| object |afe| to obtain the log-based auditory 
% spectrograms.
adsSpecTrain = transform(adsTrain,@(x)getSpectrogram(x,afe,segmentLength));
%% Define Training Labels
%%
% Use an <matlab:doc('arrayDatastore') |arrayDatastore|> to hold the training labels.
labelsTrain = arrayDatastore(adsTrain.Labels);
classes = single(unique(adsTrain.Labels));
%% Define Combined Training Datastore
% Create a combined datastore that points to the spectrogram data and the
% corresponding labels.
tdsTrain = combine(adsSpecTrain,labelsTrain);
%% Compute Validation Data
% Precompute the validation features.
adsTestT = transform(adsTest,@(x){getSpectrogram(x,afe,segmentLength)});
XTest = readall(adsTestT, UseParallel=canUseParallelPool);
XTest = cat(4,XTest{:});
%% 
% Get the validation labels.
YTest = adsTest.Labels;
%% Define Network Architecture
% Construct a small CNN as an array of layers. Use convolutional and batch
% normalization layers, and downsample the feature maps using max pooling
% layers. To reduce the possibility of overfitting, add a small amount of
% dropout to the input to the last fully connected layer.
sz = size(XTest);
specSize = sz(1:2);
imageSize = [specSize 1];
numClasses = numel(categories(YTest));
dropoutProb = 0.2;
numF = 12;
layers = [
    imageInputLayer(imageSize,Normalization="none")
    convolution2dLayer(5,numF,Padding="same")
    batchNormalizationLayer
    reluLayer
    maxPooling2dLayer(3,Stride=2,Padding="same")
    convolution2dLayer(3,2*numF,Padding="same")
    batchNormalizationLayer
    reluLayer
    maxPooling2dLayer(3,Stride=2,Padding="same")
    convolution2dLayer(3,4*numF,Padding="same")
    batchNormalizationLayer
    reluLayer
    maxPooling2dLayer(3,Stride=2,Padding="same")
    convolution2dLayer(3,4*numF,Padding="same")
    batchNormalizationLayer
    reluLayer
    convolution2dLayer(3,4*numF,Padding="same")
    batchNormalizationLayer
    reluLayer
    maxPooling2dLayer(2)
    dropoutLayer(dropoutProb)
    fullyConnectedLayer(numClasses)
    ];
%%
% To train the network with a custom training loop, convert the layers to a <matlab:doc('dlnetwork') |dlnetwork|> object.
dlnet = dlnetwork(layers);
%%
% Replace the network with your own network to train. For more information, see  <matlab:helpview('deeplearning','list_of_layers') List of Deep Learning 
% Layers>.
%% Specify Training Options
% Access the initial learning rate |initialLearnRate| defined in the
% hyperparameter table with the argument value |params.initialLearnRate|.
velocity = [];
numEpochs = 20;
initialLearnRate = params.initialLearnRate;
momentum = 0.9;
decay = 0.01;
%% 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 spectrogram data. Evaluate the model gradients and loss
% using <matlab:doc('dlfeval') |dlfeval|> and the helper function
% |modelLoss|. 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 and validation accuracy. After
% each iteration of the custom training loop, save the trained network and
% update the training progress. You can exit the custom training loop if
% you click *Stop* on the Experiment Manager toolstrip.
monitor.Metrics = ["ValidationAccuracy","TrainingLoss"];
%%
% Define a <matlab:doc('minibatchqueue') |minibatchqueue|> to create a
% mini-batch queue for the training data. Set OutputEnvironment to "auto"
% to use a GPU if it is available.
miniBatchSize = 128;
mqueue = minibatchqueue(tdsTrain,MiniBatchSize=miniBatchSize,...
                        OutputEnvironment="auto",...
                        OutputAsDlarray=false,...
                        PartialMiniBatch="discard");
%%
% Define the training loop.
iteration = 0;
output.trainingInfo.loss = [];
for epoch = 1:numEpochs

    shuffle(mqueue)
    while hasdata(mqueue)
        iteration = iteration + 1;
        [X,labels] = next(mqueue);
        X = reshape(X,[size(X,1) size(X,2) 1 miniBatchSize]);
        Y = zeros(numClasses, miniBatchSize, "single");
        for c = 1:numClasses
            Y(c,labels==classes(c)) = 1;
        end
        dlX = dlarray(X,"SSCB");

        [loss,gradients] = dlfeval(@modelLoss,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;
    act = forward(dlnet,dlarray(XTest,"SSCB"));
    [~,ind] = max(act,[],1);
    ind = extractdata(ind);
    valAcc = sum(double(YTest) == ind.')/numel(ind);
    recordMetrics(monitor,iteration, ValidationAccuracy = 100*valAcc);
    
    if monitor.Stop
        return;
    end
end
%% 
% End of training function.
end
%% Supporting Functions
% Helper function to evaluate the model loss and gradients.
function [loss,gradients] = modelLoss(dlnet,dlX,T)
dlYPred = forward(dlnet,dlX);
dlYPred = softmax(dlYPred);
loss = crossentropy(dlYPred,T);
gradients = dlgradient(loss,dlnet.Learnables);
end
% Get Speech Spectrograms
function X = getSpectrogram(x,afe,segmentLength)
% |getSpectrogram(x,afe,params)| computes a speech spectrogram for the
% signal |x| using the |audioFeatureExtractor| |afe|.
x = scaleAndResize(single(x),segmentLength);
spec = extract(afe,x).';
X = log10(spec + 1e-6);
end
% Scale and resize
function y = scaleAndResize(x,n)
% |scaleAndResize(x,n)| scales |x| by its maximum absolute value and forces
% its length to be |n| by trimming or zero-padding.
a = size(x,1);
if a < n
    frontPad = floor((n-a)/2);
    backPad = n - a - frontPad;
    y = [zeros(frontPad,size(x,2),like=x);x;zeros(backPad,size(x,2),like=x)];
elseif a > n
    frontTrim = floor((a-n)/2) + 1;
    backTrim = a - n - frontTrim;
    y = x(frontTrim:end-backTrim,:);
else
    y = x;
end

y = y./max(abs(y));
end
%% 
% _Copyright 2022 The MathWorks, Inc.
% 
%
