%% Train Network Exported from Deep Network Designer
% This example shows how to use a network exported from Deep Network Designer for transfer learning. After preparing the network in the app, you need to:
% Resize images.
% Specify training options.
% Train the network.
%% Resize Images for Transfer Learning
% For transfer learning, resize your images to match the input size of the pretrained network. To find the image input size of the network, in Deep Network Designer, examine the imageInputLayer. For GoogLeNet, the InputSize is 224x224.
%
% Unzip and load the images as an image datastore. This very small data set contains only 75 images in 5 classes. Divide the data into 70% for training and 30% for validation.

unzip('MerchData.zip');
imds = imageDatastore('MerchData', ...
    'IncludeSubfolders',true, ...
    'LabelSource','foldernames'); 
[imdsTrain,imdsValidation] = splitEachLabel(imds,0.7);

% If your training images are in a folder with subfolders for each class, you can create a datastore for your data by replacing MerchData with the folder location. Check the number of classes - you must prepare the network for transfer learning with the number of classes to match your data.
% Resize images in the image datastores to match the pretrained network GoogLeNet.

augimdsTrain = augmentedImageDatastore([224 224],imdsTrain);
augimdsValidation = augmentedImageDatastore([224 224],imdsValidation);

% You can also apply transformations to the images to help prevent the network from overfitting. For details, see imageDataAugmenter.
%
%% Set Training Options for Transfer Learning
% Before training, specify training options.
% For transfer learning, set InitialLearnRate to a small value to slow down learning in the transferred layers. In the app, you increased the learning rate factors for the fully connected layer to speed up learning in the new final layers. This combination of learning rate settings results in fast learning only in the new layers and slower learning in the other layers.
% Specify a small number of epochs. An epoch is a full training cycle on the entire training data set. For transfer learning, you do not need to train for as many epochs. Shuffle the data every epoch.
% Specify the mini-batch size, that is, how many images to use in each iteration.
% Specify validation data and validation frequency.
% Turn on the training plot to monitor progress while you train.

options = trainingOptions('sgdm', ...
    'MiniBatchSize',10, ...
    'MaxEpochs',6, ...
    'InitialLearnRate',1e-4, ...
    'Shuffle','every-epoch', ...
    'ValidationData',augimdsValidation, ...
    'ValidationFrequency',6, ...
    'Verbose',false, ...
    'Plots','training-progress');

%
%% Train Network
% To train the network, supply the layers you exported from the app, here named lgraph_1, your resized images, and training options, to the trainNetwork function. By default, trainNetwork uses a GPU if available (requires Parallel Computing Toolbox™). Otherwise, it uses a CPU. Training is fast because the data set is so small.

net = trainNetwork(augimdsTrain,lgraph_1,options);

%
%% Test Trained Network by Classifying Validation Images
% Use the fine-tuned network to classify the validation images, and calculate the classification accuracy.

[YPred,probs] = classify(net,augimdsValidation);
accuracy = mean(YPred == imdsValidation.Labels)

% Display four sample validation images with predicted labels and predicted probabilities.

idx = randperm(numel(imdsValidation.Files),4);
figure
for i = 1:4
    subplot(2,2,i)
    I = readimage(imdsValidation,idx(i));
    imshow(I)
    label = YPred(idx(i));
    title(string(label) + ", " + num2str(100*max(probs(idx(i),:)),3) + "%");
end

% Copyright 2019 The MathWorks, Inc.