% This example shows how to create and train a simple convolutional neural
% network for deep learning classification. Convolutional neural networks
% are essential tools for deep learning, and are especially suited for
% image recognition.

% -The example demonstrates how to:
% -Load and explore image data.
% -Define the network architecture.
% -Specify training options.
% -Train the network.
% -Predict the labels of new data and calculate the classification accuracy.

%% first the parameters
MaxEpochs = 4;
InitialLearnRate = 0.01;
MiniBatchSize = 75;
ratioTrainFiles = 0.75;
% flags
verbose = true;
dataAugmentation = true;
doTrain = true;
dropOut = false;
transferLearn = false;

%% Load and Explore Image Data
% Load the digit sample data as an image datastore. imageDatastore
% automatically labels the images based on folder names and stores the data
% as an ImageDatastore object. An image datastore enables you to store
% large image data, including data that does not fit in memory, and
% efficiently read batches of images during training of a convolutional
% neural network.
digitDatasetPath = fullfile(matlabroot,'toolbox','nnet','nndemos', ...
    'nndatasets','DigitDataset');
imds = imageDatastore(digitDatasetPath, ...
    'IncludeSubfolders',true,'LabelSource','foldernames');

%% Specify Training and Validation Sets
% Divide the data into training and validation data sets, so that each
% category in the training set contains 750 images, and the validation set
% contains the remaining images from each label. splitEachLabel splits the
% datastore digitData into two new datastores, trainDigitData and
% valDigitData.
labelCount = countEachLabel(imds);
[imdsTrain,imdsValidation] = splitEachLabel(imds,ratioTrainFiles*min(labelCount.Count),'randomize');
inputImagesize = size(imread(imdsTrain.Files{1}));

%% Data Augmentation
% Create an image data augmenter that preprocesses images before training.
% This augmenter rotates images by random angles in the range [0, 360]
% degrees and resizes images by random scale factors in the range [0.5, 1].
if dataAugmentation
    augmenter = imageDataAugmenter('RandRotation',[0 360],'RandScale',[0.5 1]); 
    imdsTrain = augmentedImageDatastore(inputImagesize,imdsTrain,'DataAugmentation',augmenter);
end

% Display some of the images in the datastore.
if verbose
    figure;
    perm = randperm(length(imdsTrain.Files),20);
    for i = 1:20
        subplot(4,5,i);
        imshow(imdsTrain.Files{perm(i)});
    end
end

%% Define Network Architecture
% Define the convolutional neural network architecture.
mydropOutLayer = [];
if dropOut
    mydropOutLayer = dropoutLayer(0.5)
end
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)
    mydropOutLayer
    softmaxLayer
    classificationLayer];

%% Tranfer Learning
if transferLearn
    % at this point you should have a pretrained network to transfer. how
    % transfer learning works will depend on what you will transfer.
    % steps:
    % 1. load the pretrained network, 
    %    (e.g.) I have loaded "myPreTrainedNetwork" to my workspace
    % 2. Transfer necessary layers layers(2) (be careful must be compatible!)
    %    (e.g.) net.Layers(i).Weights = myPreTrainedNetwork.Layers(j).Weights
    %    (e.g.) net.Layers(i).Bias = myPreTrainedNetwork.Layers(j).Bias
    % 3. freeze or reassign learning rates for the transfered layers
    %    (e.g.) net.Layers(i).WeightLearnRateFactor = alfa
    %    this means that the learning rate for this layers alfa times the
    %    global learning rate (which is initially and globally --> "InitialLearnRate")
    %    make alfa zero, that layer will "freeze" during training.
end

%% Specify Training Options
% After defining the network structure, specify the training options. Train
% the network using stochastic gradient descent with momentum (SGDM) with
% an initial learning rate of 0.01. Set the maximum number of epochs to 4.
% An epoch is a full training cycle on the entire training data set.
% Monitor the network accuracy during training by specifying validation
% data and validation frequency. Shuffle the data every epoch. The software
% trains the network on the training data and calculates the accuracy on
% the validation data at regular intervals during training. The validation
% data is not used to update the network weights. Turn on the training
% progress plot, and turn off the command window output.
options = trainingOptions('sgdm', ...
    'InitialLearnRate',0.01, ...
    'MaxEpochs',MaxEpochs, ...
    'Shuffle','every-epoch', ...
    'ValidationData',imdsValidation, ...
    'ValidationFrequency',30, ...
    'Verbose',false, ...
    'Plots','training-progress',...
    'MiniBatchSize',MiniBatchSize);

%% Train Network Using Training Data
% Train the network using the architecture defined by layers, the training
% data, and the training options. By default, trainNetwork uses a GPU if
% one is available (requires Parallel Computing Toolbox™ and a CUDA®
% enabled GPU with compute capability 3.0 or higher). Otherwise, it uses a
% CPU. You can also specify the execution environment by using the
% 'ExecutionEnvironment' name-value pair argument of trainingOptions.

% The training progress plot shows the mini-batch loss and accuracy and the
% validation loss and accuracy. For more information on the training
% progress plot, see Monitor Deep Learning Training Progress. The loss is
% the cross-entropy loss. The accuracy is the percentage of images that the
% network classifies correctly.
if doTrain
    net = trainNetwork(imdsTrain,layers,options);
end