%% Image Generation Using Generative Adversarial Network (GAN)
% 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 generating 
% images using a GAN, see <matlab:helpview('deeplearning','TrainGenerativeAdversarialNetworkGANExample') 
% Train Generative Adversarial Network (GAN)>. 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
%% 
% * |output| is a structure with two fields, |dlnetGenerator| and |dlnetDiscriminator|, 
% containing the final trained networks.
function output = {functionName}(params,monitor)
%% Initialize Output
output = [];
%% Load Training Data
% Load the training data as an <matlab:doc('imageDatastore') |imageDatastore|>. 
% The training data consists of 3670 images from the <https://www.tensorflow.org/datasets/catalog/tf_flowers 
% Flowers> data set. Augment the data to include random horizontal flipping and 
% resize the images to have size 64-by-64.
monitor.Status = "Loading Data";
url = 'http://download.tensorflow.org/example_images/flower_photos.tgz';
downloadFolder = tempdir;
filename = fullfile(downloadFolder,'flower_dataset.tgz');
imageFolder = fullfile(downloadFolder,'flower_photos');
if ~exist(imageFolder,'dir')
    disp('Downloading Flowers data set (218 MB)...')
    websave(filename,url);
    untar(filename,downloadFolder)
end
datasetFolder = fullfile(imageFolder);
imdsTrain = imageDatastore(datasetFolder, ...
    'IncludeSubfolders',true);
augmenter = imageDataAugmenter('RandXReflection',true);
augimdsTrain = augmentedImageDatastore([64 64],imdsTrain, ...
    'DataAugmentation',augmenter);
%% 
% 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 Generator Network
% Define the architecture for the generator network as a layer graph that generates 
% images from 1-by-1-by-100 arrays of random values. To train the network with 
% a custom training loop and enable automatic differentiation, convert the layer 
% graph to a <matlab:doc('dlnetwork') |dlnetwork|> object.
monitor.Status = "Creating Networks";
filterSize = 5;
numFilters = 64;
numLatentInputs = 100;
projectionSize = [4 4 512];
addpath(fullfile(matlabroot,'examples','nnet','main'))
layersGenerator = [
    imageInputLayer([1 1 numLatentInputs],'Normalization','none', ...
        'Name','in')
    projectAndReshapeLayer(projectionSize,numLatentInputs,'Name','proj');
    transposedConv2dLayer(filterSize,4*numFilters,'Name','tconv1')
    batchNormalizationLayer('Name','bnorm1')
    reluLayer('Name','relu1')
    transposedConv2dLayer(filterSize,2*numFilters,'Stride',2, ...
        'Cropping','same','Name','tconv2')
    batchNormalizationLayer('Name','bnorm2')
    reluLayer('Name','relu2')
    transposedConv2dLayer(filterSize,numFilters,'Stride',2, ...
        'Cropping','same','Name','tconv3')
    batchNormalizationLayer('Name','bnorm3')
    reluLayer('Name','relu3')
    transposedConv2dLayer(filterSize,3,'Stride',2, ...
        'Cropping','same','Name','tconv4')
    tanhLayer('Name','tanh')];
lgraphGenerator = layerGraph(layersGenerator);
output.dlnetGenerator = dlnetwork(lgraphGenerator);
%% 
% 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>.
%% Define Discriminator Network
% Define the architecture for the discriminator network as a layer graph that 
% classifies real and generated 64-by-64-by-3 images. To train the network with 
% a custom training loop and enable automatic differentiation, convert the layer 
% graph to a <matlab:doc('dlnetwork') |dlnetwork|> object.
dropoutProb = 0.5;
numFilters = 64;
scale = 0.2;
inputSize = [64 64 3];
filterSize = 5;
layersDiscriminator = [
    imageInputLayer(inputSize,'Normalization','none','Name','in')
    dropoutLayer(dropoutProb,'Name','dropout')
    convolution2dLayer(filterSize,numFilters,'Stride',2, ...
        'Padding','same','Name','conv1')
    leakyReluLayer(scale,'Name','lrelu1')
    convolution2dLayer(filterSize,2*numFilters,'Stride',2, ...
        'Padding','same','Name','conv2')
    batchNormalizationLayer('Name','bn2')
    leakyReluLayer(scale,'Name','lrelu2')
    convolution2dLayer(filterSize,4*numFilters,'Stride',2, ...
        'Padding','same','Name','conv3')
    batchNormalizationLayer('Name','bn3')
    leakyReluLayer(scale,'Name','lrelu3')
    convolution2dLayer(filterSize,8*numFilters,'Stride',2, ...
        'Padding','same','Name','conv4')
    batchNormalizationLayer('Name','bn4')
    leakyReluLayer(scale,'Name','lrelu4')
    convolution2dLayer(4,1,'Name','conv5')];
lgraphDiscriminator = layerGraph(layersDiscriminator);
output.dlnetDiscriminator = dlnetwork(lgraphDiscriminator);
%% 
% 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 128 for 10 epochs. For both networks, specify 
% a gradient decay factor of 0.5 and a squared gradient decay factor of 0.999. 
% Initialize the trailing average gradient and trailing average gradient-square 
% decay rates as |[]|. If the discriminator learns to discriminate between real 
% and generated images too quickly, then the generator can fail to train. To better 
% balance the learning of the discriminator and the generator, add noise to the 
% real data by randomly flipping the labels. Specify to flip 30% of the real labels. 
% Train on a GPU if one is available.
% 
% Training can take some time to run. For better results, consider increasing 
% the training to 500 epochs. For larger data sets, you might not need to train 
% for as many epochs. 
numEpochs = 10;
miniBatchSize = 128;
gradientDecayFactor = 0.5;
squaredGradientDecayFactor = 0.999;
trailingAvgGenerator = [];
trailingAvgSqGenerator = [];
trailingAvgDiscriminator = [];
trailingAvgSqDiscriminator = [];
flipFactor = 0.3;
executionEnvironment = "auto";
%% Train Model
% Use the input argument |monitor| to track the progress of the training and 
% display information (|gradientDecayFactor| and |squaredGradientDecayFactor|) 
% and metric values (|scoreGenerator| and |scoreDiscriminator|) in the experiment 
% results table.
monitor.Info = ["gradientDecayFactor","squaredGradientDecayFactor"];
monitor.Metrics = ["scoreGenerator","scoreDiscriminator"];
monitor.XLabel = "Iteration";
groupSubPlot(monitor,"Score",monitor.Metrics);
updateInfo(monitor, ...
    'gradientDecayFactor',gradientDecayFactor, ...
    'squaredGradientDecayFactor',squaredGradientDecayFactor);
monitor.Status = "Training";
%% 
% Use <matlab:doc('minibatchqueue') |minibatchqueue|> to process and manage 
% the mini-batches of images. For each mini-batch:
%% 
% * Use the mini-batch preprocessing helper function |preprocessMiniBatch| to 
% rescale the images in the range [-1,1].
% * Discard any partial mini-batches with less than 128 observations.
% * Format the image data with the dimension labels |'SSCB'| (spatial, spatial, 
% channel, batch). By default, the |minibatchqueue| object converts the data to 
% |dlarray| objects with underlying type |single|.
augimdsTrain.MiniBatchSize = miniBatchSize;
mbq = minibatchqueue(augimdsTrain,...
    'MiniBatchSize',miniBatchSize,...
    'PartialMiniBatch','discard',...
    'MiniBatchFcn', @preprocessMiniBatch,...
    'MiniBatchFormat','SSCB',...
    'OutputEnvironment',executionEnvironment);
%% 
% Train the model using a custom training loop. For each epoch, shuffle the 
% datastore and loop over mini-batches of data. If training on a GPU, then convert 
% latent inputs to <matlab:doc('gpuArray') |gpuArray|>. Evaluate the model gradients 
% and update the discriminator and generator network parameters. After each iteration 
% of the custom training loop, save the trained network, record the metric values, 
% and update the training progress. Exit the custom training loop if you click 
% *Stop* in the Experiment Manager toolstrip. 
iteration = 0;
for epoch = 1:numEpochs
    shuffle(mbq);
    while hasdata(mbq)
        iteration = iteration + 1;
        dlX = next(mbq);
        
        Z = randn(1,1,numLatentInputs,size(dlX,4),'single');
        dlZ = dlarray(Z,'SSCB');
        
        if (executionEnvironment == "auto" && canUseGPU) || executionEnvironment == "gpu"
            dlZ = gpuArray(dlZ);
        end
        
        [gradientsGenerator,gradientsDiscriminator,stateGenerator,scoreGenerator,scoreDiscriminator] = ...
            dlfeval(@modelGradients, output.dlnetGenerator, output.dlnetDiscriminator, dlX, dlZ, flipFactor);
        output.dlnetGenerator.State = stateGenerator;
        
        [output.dlnetDiscriminator,trailingAvgDiscriminator,trailingAvgSqDiscriminator] = ...
            adamupdate(output.dlnetDiscriminator, gradientsDiscriminator, ...
            trailingAvgDiscriminator, trailingAvgSqDiscriminator, iteration, ...
            initialLearnRate, gradientDecayFactor, squaredGradientDecayFactor);
        
        [output.dlnetGenerator,trailingAvgGenerator,trailingAvgSqGenerator] = ...
            adamupdate(output.dlnetGenerator, gradientsGenerator, ...
            trailingAvgGenerator, trailingAvgSqGenerator, iteration, ...
            initialLearnRate, gradientDecayFactor, squaredGradientDecayFactor);
        
        scoreGeneratorValue = ...
            double(gather(extractdata(scoreGenerator)));
        scoreDiscriminatorValue = ...
            double(gather(extractdata(scoreDiscriminator)));
        recordMetrics(monitor,iteration, ...
            'scoreGenerator',scoreGeneratorValue, ...
            'scoreDiscriminator',scoreDiscriminatorValue);
        
        if monitor.Stop
            return;
        end
    end
    monitor.Progress = (epoch/numEpochs)*100;
end
%% 
% End of training function.
end
%% Supporting Functions for Training
% This function preprocesses the data using the following steps:
%% 
% # Extract the image data from the incoming cell array and concatenate into 
% a numeric array.
% # Rescale the images to be in the range [-1,1].
function X = preprocessMiniBatch(data)
    X = cat(4,data{:});
    X = rescale(X,-1,1,'InputMin',0,'InputMax',255);
end
%% 
% This function takes as input the generator and discriminator |dlnetwork| objects 
% (|dlnetGenerator| and |dlnetDiscriminator|), a mini-batch of input data (|dlX|), 
% an array of random values (|dlZ|), and the percentage of real labels to flip 
% (|flipFactor|), and returns the gradients of the loss with respect to the learnable 
% parameters in the networks, the generator state, and the scores of the two networks. 
% Because the discriminator output is not in the range [0,1], |modelGradients| 
% applies the |sigmoid| function to convert this value into a probability.
function [gradientsGenerator,gradientsDiscriminator,stateGenerator,scoreGenerator,scoreDiscriminator] = ...
    modelGradients(dlnetGenerator,dlnetDiscriminator,dlX,dlZ,flipFactor)
    dlYPred = forward(dlnetDiscriminator,dlX);
    [dlXGenerated,stateGenerator] = forward(dlnetGenerator,dlZ);
    dlYPredGenerated = forward(dlnetDiscriminator, dlXGenerated);
    probGenerated = sigmoid(dlYPredGenerated);
    probReal = sigmoid(dlYPred);
    scoreDiscriminator = ((mean(probReal)+mean(1-probGenerated))/2);
    scoreGenerator = mean(probGenerated);
    numObservations = size(probReal,4);
    idx = randperm(numObservations,floor(flipFactor*numObservations));
    probReal(:,:,:,idx) = 1-probReal(:,:,:,idx);
    [lossGenerator, lossDiscriminator] = GANLoss(probReal,probGenerated);
    gradientsGenerator = dlgradient(lossGenerator, ...
        dlnetGenerator.Learnables,'RetainData',true);
    gradientsDiscriminator = dlgradient(lossDiscriminator, ...
        dlnetDiscriminator.Learnables);
end
%% 
% This function returns the loss for the discriminator and generator networks.
function [lossGenerator,lossDiscriminator] = ...
    GANLoss(probReal,probGenerated)
    lossDiscriminator =  -mean(log(probReal)) -mean(log(1-probGenerated));
    lossGenerator = -mean(log(probGenerated));
end
%% 
% Copyright 2020-2021 The MathWorks, Inc.