%% Image Comparison Using Siamese Network
% 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 comparing images 
% using a Siamese network, see <matlab:helpview('deeplearning','TrainASiameseNetworkToCompareImagesExample') 
% Train a Siamese Network to Compare Images>. For more information on custom training 
% 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
%% 
% * |dlnet| is the final trained subnetwork.
function dlnet = {functionName}(params,monitor)
%% Initialize Output
dlnet = [];
%% Load and Preprocess Training Data
% Load the training data as an <matlab:doc('imageDatastore') |imageDatastore|>. 
% The training data consists of images from the <https://github.com/brendenlake/omniglot 
% Omniglot> training data set. Specify the labels manually by extracting the labels 
% from the file names.
monitor.Status = "Loading Data";
url = "https://github.com/brendenlake/omniglot/raw/master/python/images_background.zip";
downloadFolder = tempdir;
filename = fullfile(downloadFolder,"images_background.zip");
dataFolderTrain = fullfile(downloadFolder,'images_background');
if ~exist(dataFolderTrain,"dir")
    monitor.Status = "Loading Omniglot Data";
    websave(filename,url);
    unzip(filename,downloadFolder);
end
imdsTrain = imageDatastore(dataFolderTrain, ...
    'IncludeSubfolders',true, ...
    'LabelSource','none');
files = imdsTrain.Files;
parts = split(files,filesep);
labels = join(parts(:,(end-2):(end-1)),'_');
imdsTrain.Labels = categorical(labels);
%% 
% 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
% Create two identical subnetworks that accept 105-by-105-by-1 images and output 
% a feature vector. To train the network with a custom training loop and enable 
% automatic differentiation, convert the layer graph to a <matlab:doc('dlnetwork') 
% |dlnetwork|> object. Create the weights for the final |fullyconnect| operation. 
% Initialize the weights by sampling a random selection from a narrow normal distribution 
% with standard deviation of 0.01.
monitor.Status = "Creating Network";
layers = [
    imageInputLayer([105 105 1],'Name','input1','Normalization','none')
    convolution2dLayer(10,64,'Name','conv1', ...
        'WeightsInitializer','narrow-normal', ...
        'BiasInitializer','narrow-normal')
    reluLayer('Name','relu1')
    maxPooling2dLayer(2,'Stride',2,'Name','maxpool1')
    convolution2dLayer(7,128,'Name','conv2', ...
        'WeightsInitializer','narrow-normal', ...
        'BiasInitializer','narrow-normal')
    reluLayer('Name','relu2')
    maxPooling2dLayer(2,'Stride',2,'Name','maxpool2')
    convolution2dLayer(4,128,'Name','conv3', ...
        'WeightsInitializer','narrow-normal', ...
        'BiasInitializer','narrow-normal')
    reluLayer('Name','relu3')
    maxPooling2dLayer(2,'Stride',2,'Name','maxpool3')
    convolution2dLayer(5,256,'Name','conv4', ...
        'WeightsInitializer','narrow-normal', ...
        'BiasInitializer','narrow-normal')
    reluLayer('Name','relu4')
    fullyConnectedLayer(4096,'Name','fc1', ...
        'WeightsInitializer','narrow-normal', ...
        'BiasInitializer','narrow-normal')];
lgraph = layerGraph(layers);
dlnet = dlnetwork(lgraph);
fcWeights = dlarray(0.01*randn(1,4096));
fcBias = dlarray(0.01*randn(1,1));
fcParams = struct(...
    "FcWeights",fcWeights,...
    "FcBias",fcBias);
%% 
% 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|.
initialLearnRate = params.myInitialLearnRate;
%% 
% Train with a mini-batch size of 180 for 100 iterations. Set the gradient decay 
% factor to 0.9 and the squared gradient decay factor to 0.99. Initialize the 
% trailing average gradient and trailing average gradient-square decay rates as 
% |[]|. Train on a GPU if one is available.
% 
% Training can take some time to run. For better results, consider increasing 
% the training to 10,000 iterations.
numIterations = 100;
miniBatchSize = 180;
gradientDecayFactor = 0.9;
squaredGradientDecayFactor = 0.99;
trailingAvgSubnet = [];
trailingAvgSqSubnet = [];
trailingAvgParams = [];
trailingAvgSqParams = [];
executionEnvironment = "auto";
%% Train Model
% Use the input argument |monitor| to track the progress of the training and 
% display information (|gradientDecayFactor| and |squaredGradientDecayFactor|) 
% and training loss in the experiment results table.
monitor.Info = ["gradientDecayFactor","squaredGradientDecayFactor"];
monitor.Metrics = "TrainingLoss";
monitor.XLabel = "Iteration";
updateInfo(monitor, ...
    'gradientDecayFactor',gradientDecayFactor, ...
    'squaredGradientDecayFactor',squaredGradientDecayFactor);
monitor.Status = "Training";
%% 
% Train the model using a custom training loop. Loop over the training data 
% and update the network parameters at each iteration. Extract a batch of image 
% pairs and labels using the |getSiameseBatch| helper function. 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 and |'CB'| (channel, batch) for the labels. If training on a 
% GPU, then convert the data to <matlab:doc('gpuArray') |gpuArray|>. Evaluate 
% the model gradients and update the network parameters. After each iteration 
% of the custom training loop, save the trained network, record the training loss, 
% and update the training progress. Exit the custom training loop if you click 
% *Stop* in the Experiment Manager toolstrip. 
for iteration = 1:numIterations
    [X1,X2,pairLabels] = getSiameseBatch(imdsTrain,miniBatchSize);
    
    dlX1 = dlarray(single(X1),'SSCB');
    dlX2 = dlarray(single(X2),'SSCB');
    
    if (executionEnvironment == "auto" && canUseGPU) || executionEnvironment == "gpu"
        dlX1 = gpuArray(dlX1);
        dlX2 = gpuArray(dlX2);
    end
    
    [gradientsSubnet, gradientsParams,loss] = dlfeval(@modelGradients, ...
        dlnet,fcParams,dlX1,dlX2,pairLabels);
    lossValue = double(gather(extractdata(loss)));
    
    [dlnet,trailingAvgSubnet,trailingAvgSqSubnet] = adamupdate(dlnet, ...
        gradientsSubnet,trailingAvgSubnet,trailingAvgSqSubnet, ...
        iteration,initialLearnRate,gradientDecayFactor,squaredGradientDecayFactor);
    
    [fcParams,trailingAvgParams,trailingAvgSqParams] = ...
        adamupdate(fcParams,gradientsParams, ...
        trailingAvgParams,trailingAvgSqParams,iteration,initialLearnRate,gradientDecayFactor,squaredGradientDecayFactor);
    recordMetrics(monitor, iteration, 'TrainingLoss', lossValue);
    monitor.Progress = (iteration/numIterations)*100;
    if monitor.Stop
        return;
    end
end
%% 
% End of training function.
end
%% Supporting Functions for Training
% This function takes as input the Siamese |dlnetwork| object |net|, a pair 
% of mini-batch input data |dlX1| and |dlX2|, and the label indicating whether 
% they are similar or dissimilar. The function returns the gradients of the loss 
% with respect to the learnable parameters in the network and the binary cross-entropy 
% loss between the prediction and the ground truth.
function [gradientsSubnet,gradientsParams,loss] = modelGradients(dlnet,fcParams,dlX1,dlX2,pairLabels)
    Y = forwardSiamese(dlnet,fcParams,dlX1,dlX2);
    loss = binarycrossentropy(Y,pairLabels);
    [gradientsSubnet,gradientsParams] = dlgradient(loss,dlnet.Learnables,fcParams);
end
%% 
% This function defines how the subnetworks and the |fullyconnect| and |sigmoid| 
% operations combine to form the complete Siamese network. The function accepts 
% the network structure and two training images and returns a prediction of the 
% probability of the pair being similar (closer to 1) or dissimilar (closer to 
% 0).
function Y = forwardSiamese(dlnet,fcParams,dlX1,dlX2)
    F1 = forward(dlnet,dlX1);
    F1 = sigmoid(F1);
    F2 = forward(dlnet,dlX2);
    F2 = sigmoid(F2);
    Y = abs(F1 - F2);
    Y = fullyconnect(Y,fcParams.FcWeights,fcParams.FcBias);
    Y = sigmoid(Y);
end
%% 
% This function returns the binary cross-entropy loss value for a prediction 
% from the network.
function loss = binarycrossentropy(Y,pairLabels)
    precision = underlyingType(Y);
    Y(Y < eps(precision)) = eps(precision);
    Y(Y > 1 - eps(precision)) = 1 - eps(precision);
    loss = -pairLabels.*log(Y) - (1 - pairLabels).*log(1 - Y);
    loss = sum(loss)/numel(pairLabels);
end
%% 
% This function returns a randomly selected batch of paired images. On average, 
% this function produces a balanced set of similar and dissimilar pairs.
function [X1,X2,pairLabels] = getSiameseBatch(imds,miniBatchSize)
    pairLabels = zeros(1,miniBatchSize);
    imgSize = size(readimage(imds,1));
    X1 = zeros([imgSize 1 miniBatchSize]);
    X2 = zeros([imgSize 1 miniBatchSize]);
    for i = 1:miniBatchSize
        choice = rand(1);
        if choice < 0.5
            [pairIdx1,pairIdx2,pairLabels(i)] = getSimilarPair(imds.Labels);
        else
            [pairIdx1,pairIdx2,pairLabels(i)] = getDissimilarPair(imds.Labels);
        end
        X1(:,:,:,i) = imds.readimage(pairIdx1);
        X2(:,:,:,i) = imds.readimage(pairIdx2);
    end
end
%% 
% This function returns a random pair of indices for images that are in the 
% same class and the similar pair label of 1.
function [pairIdx1,pairIdx2,label] = getSimilarPair(classLabel)
    classes = unique(classLabel);
    classChoice = randi(numel(classes));
    idxs = find(classLabel==classes(classChoice));
    pairIdxChoice = randperm(numel(idxs),2);
    pairIdx1 = idxs(pairIdxChoice(1));
    pairIdx2 = idxs(pairIdxChoice(2));
    label = 1;
end
%% 
% This function returns a random pair of indices for images that are in different 
% classes and the dissimilar pair label of 0.
function  [pairIdx1,pairIdx2,label] = getDissimilarPair(classLabel)
    classes = unique(classLabel);
    classesChoice = randperm(numel(classes),2);
    idxs1 = find(classLabel==classes(classesChoice(1)));
    idxs2 = find(classLabel==classes(classesChoice(2)));
    pairIdx1Choice = randi(numel(idxs1));
    pairIdx2Choice = randi(numel(idxs2));
    pairIdx1 = idxs1(pairIdx1Choice);
    pairIdx2 = idxs2(pairIdx2Choice);
    label = 0;
end
%% 
% Copyright 2020-2021 The MathWorks, Inc.