%% Use Experiment Manager to Train Generative Adversarial Networks (GANs)
% This example shows how to create a custom training experiment to train a
% generative adversarial network (GAN) that generates images of flowers.
% For a custom training experiment, you explicitly define the training
% procedure used by
% *<docid:nnet_ref#mw_11ce840f-eed7-4975-9c28-baeebd79f092 Experiment
% Manager>*. In this example, you implement a custom training loop to train
% a GAN, a type of deep learning network that can generate data with
% similar characteristics as the input real data. A GAN consists of two
% networks that train together:
% 
% * Generator &mdash; Given a vector of random values (latent inputs) as
% input, this network generates data with the same structure as the
% training data.
% * Discriminator &mdash; Given batches of data containing observations
% from both the training data, and generated data from the generator, this
% network attempts to classify the observations as "real" or "generated."
%
% <<../exp_mgr_gan_net.png>>
%
% To train a GAN, train both networks simultaneously to maximize the
% performance of both networks:
% 
% * Train the generator to generate data that "fools" the discriminator. To
% optimize the performance of the generator, maximize the loss of the
% discriminator when given generated data. In other words, the objective of
% the generator is to generate data that the discriminator classifies as
% "real."
% * Train the discriminator to distinguish between real and generated data.
% To optimize the performance of the discriminator, minimize the loss of
% the discriminator when given batches of both real and generated data. In
% other words, the objective of the discriminator is to not be "fooled" by
% the generator.
% 
% Ideally, these strategies result in a generator that generates
% convincingly realistic data and a discriminator that has learned strong
% feature representations that are characteristic of the training data. For
% more information, see
% <docid:nnet_ug#mw_eb769257-5822-4140-b41d-c6c19ffa4727 Train Generative
% Adversarial Network (GAN)>.

% Copyright 2020-2021 The MathWorks, Inc.

%% Open Experiment
% First, open the example. Experiment Manager loads a project with a
% preconfigured experiment that you can inspect and run. To open the
% experiment, in the *Experiment Browser* pane, double-click the name of
% the experiment (|ImageGenerationExperiment|).
%
% <<../exp_mgr_gan_config.png>>

% Project Name: 'FlowerImageGenerationProject'

%% 
% Custom training experiments consist of a description, a table of
% hyperparameters, and a training function. For more information, see
% <docid:nnet_ref#mw_b6368d56-fbf4-4a60-8ddb-6404424c570c
% Configure Custom Training Experiment>.
%
% The *Description* field contains a textual description of the experiment.
% For this example, the description is:
%
%  Train a generative adversarial network (GAN) to generate images of flowers.
%  Use hyperparameters to specify:
%  * the probability of the dropout layer in the discriminator network
%  * the fraction of real labels to flip while training the discriminator network
%
% The *Hyperparameters* section specifies the strategy (|Exhaustive Sweep|)
% and hyperparameter values to use for the experiment. When you run the
% experiment, Experiment Manager trains the network using every combination
% of hyperparameter values specified in the hyperparameter table. This
% example uses two hyperparameters:
%
% * |dropoutProb| sets the probability of the dropout layer in the
% discriminator network. By default, the values for this hyperparameter are
% specified as |[0.25 0.5 0.75]|.
% * |flipFactor| sets the fraction of real labels to flip when you train
% the discriminator network. The experiment uses this hyperparameter to add
% noise to the real data and better balance the learning of the
% discriminator and the generator. Otherwise, if the discriminator learns
% to discriminate between real and generated images too quickly, then the
% generator can fail to train. The values for this hyperparameter are
% specified as |[0.1 0.3 0.5]|.
%
% The *Training Function* specifies the training data, network
% architecture, training options, and training procedure used by the
% experiment. The input to the training function is a structure with fields
% from the hyperparameter table and an <docid:nnet_ref#experiments_monitor
% experiments.Monitor> object that you can use to track the progress of the
% training, record values of the metrics used by the training, and produce
% training plots. The training function returns a structure that contains
% the trained generator network, the trained discriminator network, and the
% execution environment used for training. Experiment Manager saves this
% output, so you can export it to the MATLAB workspace when the training is
% complete. The training function has six sections.
% 
% * *Initialize Output* sets the initial value of the networks to empty
% arrays to indicate that the training has not started. The experiment sets
% the execution environment to |"auto"|, so it trains the networks on a GPU
% if one is available. Using a GPU requires Parallel Computing Toolbox(TM)
% and a supported GPU device. For more information, see
% <docid:distcomp_ug#mw_57e04559-0b60-42d5-ad55-e77ec5f5865f GPU Support by
% Release>.
%   
%   output.generator = [];
%   output.discriminator = [];
%   output.executionEnvironment = "auto";
%
% * *Load Training Data* defines the training data for the experiment as an
% |imageDatastore| object. The experiment uses the Flowers data set, which
% contains 3670 images of flowers and is about 218 MB. For more information
% on this data set, see
% <docid:nnet_ug#mw_d3130147-29e6-47eb-b6f3-9a11a4c8b659 Image Data Sets>.
%   
%   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")
%       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);
%
% * *Define Generator Network* defines 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, the training function converts the
% layer graph to a <docid:nnet_ref#object_dlnetwork dlnetwork> object.
%   
%   monitor.Status = "Creating Generator";
%
%   filterSize = 5;
%   numFilters = 64;
%   numLatentInputs = 100;
%   projectionSize = [4 4 512];
%
%   layersGenerator = [
%       featureInputLayer(numLatentInputs,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.generator = dlnetwork(lgraphGenerator);
%
% * *Define Discriminator Network* defines the architecture for the
% discriminator network as a layer graph that classifies real and generated
% 64-by-64-by-3 images. The dropout layer uses the dropout probability
% defined in the hyperparameter table. To train the network with a custom
% training loop and enable automatic differentiation, the training function
% converts the layer graph to a <docid:nnet_ref#object_dlnetwork dlnetwork>
% object.
%   
%   monitor.Status = "Creating Discriminator";
%
%   filterSize = 5;
%   numFilters = 64;
%   inputSize = [64 64 3];
%   dropoutProb = params.dropoutProb;
%   scale = 0.2;
%
%   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.discriminator = dlnetwork(lgraphDiscriminator);
%
% * *Specify Training Options* defines the training options used by the
% experiment. In this example, Experiment Manager trains the networks with
% a mini-batch size of 128 for 50 epochs using an initial learning rate of
% 0.0002, a gradient decay factor of 0.5, and a squared gradient decay
% factor of 0.999.
%   
%   numEpochs = 50;
%   miniBatchSize = 128;
%   initialLearnRate = 0.0002;
%   gradientDecayFactor = 0.5;
%   squaredGradientDecayFactor = 0.999;
%   trailingAvgGenerator = [];
%   trailingAvgSqGenerator = [];
%   trailingAvgDiscriminator = [];
%   trailingAvgSqDiscriminator = [];
%   flipFactor = params.flipFactor;
%
% * *Train Model* defines the custom training loop used by the experiment.
% The custom training loop uses
% <docid:nnet_ref#mw_412afe49-5c0e-4f1b-8df7-b229f1b07eec minibatchqueue>
% to process and manage the mini-batches of images. For each mini-batch,
% the |minibatchqueue| object rescales the images in the range [-1,1],
% discards any partial mini-batches with fewer than 128 observations, and
% formats the image data with the dimension labels |'SSCB'| (spatial,
% spatial, channel, batch). By default, the |minibatchqueue| object
% converts the data to
% <docid:nnet_ref#mw_bc7bf07e-0207-40d7-8568-5bdd002c6390 dlarray> objects
% with underlying type |single|. For each epoch, the custom training loop
% shuffles the datastore and loops over mini-batches of data. If you train
% on a GPU, the data is converted to <docid:distcomp_ug#bsni4_e-1 gpuArray>
% objects. Then, the training function evaluates the model gradients and
% updates the discriminator and generator network parameters. After each
% iteration of the custom training loop, the training function saves the
% trained networks and updates the training progress.
%   
%   monitor.Metrics = ["scoreGenerator","scoreDiscriminator","scoreCombined"];
%   monitor.XLabel = "Iteration";
%   groupSubPlot(monitor,"Combined Score","scoreCombined");
%   groupSubPlot(monitor,"Generator and Discriminator Scores", ...
%       ["scoreGenerator","scoreDiscriminator"]);
%   monitor.Status = "Training";
%   
%   augimdsTrain.MiniBatchSize = miniBatchSize;
%   mbq = minibatchqueue(augimdsTrain,...
%       MiniBatchSize=miniBatchSize,...
%       PartialMiniBatch="discard",...
%       MiniBatchFcn=@preprocessMiniBatch,...
%       MiniBatchFormat="SSCB",...
%       OutputEnvironment=output.executionEnvironment);
%   
%   iteration = 0;
%   for epoch = 1:numEpochs
%       shuffle(mbq);
%       while hasdata(mbq)
%           iteration = iteration + 1;
%           dlX = next(mbq);
%           Z = randn(numLatentInputs,miniBatchSize,"single");
%           dlZ = dlarray(Z,"CB");
%           if (output.executionEnvironment == "auto" && canUseGPU) || ...
%                   output.executionEnvironment == "gpu"
%               dlZ = gpuArray(dlZ);
%           end
%           [gradientsGenerator,gradientsDiscriminator,stateGenerator,scoreGenerator,scoreDiscriminator] = ...
%               dlfeval(@modelGradients,output.generator,output.discriminator,dlX,dlZ,flipFactor);
%           output.generator.State = stateGenerator;
%           [output.discriminator,trailingAvgDiscriminator,trailingAvgSqDiscriminator] = ...
%               adamupdate(output.discriminator,gradientsDiscriminator, ...
%               trailingAvgDiscriminator,trailingAvgSqDiscriminator,iteration, ...
%               initialLearnRate,gradientDecayFactor,squaredGradientDecayFactor);
%           [output.generator,trailingAvgGenerator,trailingAvgSqGenerator] = ...
%               adamupdate(output.generator,gradientsGenerator, ...
%               trailingAvgGenerator,trailingAvgSqGenerator,iteration, ...
%               initialLearnRate,gradientDecayFactor,squaredGradientDecayFactor);
%           scoreGeneratorValue = ...
%               double(gather(extractdata(scoreGenerator)));
%           scoreDiscriminatorValue = ...
%               double(gather(extractdata(scoreDiscriminator)));
%           scoreCombinedValue = 1-2*max(abs(scoreDiscriminatorValue-0.5),abs(scoreGeneratorValue-0.5));
%           recordMetrics(monitor,iteration, ...
%               scoreGenerator=scoreGeneratorValue, ...
%               scoreDiscriminator=scoreDiscriminatorValue, ...
%               scoreCombined=scoreCombinedValue);
%           if monitor.Stop || isnan(scoreGeneratorValue) || isnan(scoreDiscriminatorValue)
%               return;
%           end
%       end
%       monitor.Progress = (epoch/numEpochs)*100;
%   end
%
% Training GANs can be a challenging task because the generator and the
% discriminator networks compete against each other during the training. If
% one network learns too quickly, then the other network can fail to learn.
% To help you diagnose issues and monitor how well the generator and
% discriminator achieve their respective goals, this experiment displays a
% pair of scores in the training plot. The generator score |scoreGenerator|
% measures the likelihood that the discriminator can correctly distinguish
% generated images. The discriminator score |scoreDiscriminator| measures
% the likelihood that the discriminator can correctly distinguish all input
% images, assuming that the numbers of real and generated images passed to
% the discriminator are equal. In the ideal case, both scores are 0.5.
% Scores that are too close to zero or one can indicate that one network
% dominates the other. See
% <docid:nnet_ug#mw_1e0e9d04-f6fc-4e77-97e7-c6aab7d7e0df Monitor GAN
% Training Progress and Identify Common Failure Modes>.
%
% To help you decide which trial produces the best results, this experiment
% combines the generator score and discriminator scores into a single
% numeric value, |scoreCombined|. This metric uses the _L_-&infin; norm to
% determine how close the two networks are from the ideal scenario. It
% takes a value of one if both network scores equal 0.5, and zero if one of
% the network scores equals zero or one.
%
% To inspect the training function, under *Training Function*, click
% *Edit*. The training function opens in MATLAB(R) Editor.  In addition,
% the code for the training function appears in <#7 Appendix 1> at the end
% of this example.

%% Run Experiment
% When you run the experiment, Experiment Manager trains the network
% defined by the training function multiple times. Each trial uses a
% different combination of hyperparameter values. By default, Experiment
% Manager runs one trial at a time. If you have Parallel Computing Toolbox,
% you can run multiple trials at the same time. For best results, before
% you run your experiment, start a parallel pool with as many workers as
% GPUs. For more information, see
% <docid:nnet_ug#mw_4858679a-7447-4736-9c57-7b8dc706648d Use Experiment
% Manager to Train Networks in Parallel>.
%
% * To run one trial of the experiment at a time, on the Experiment Manager
% toolstrip, click *Run*.
% * To run multiple trials at the same time, click *Use Parallel* and then
% *Run*. If there is no current parallel pool, Experiment Manager starts
% one using the default cluster profile. Experiment Manager then executes
% multiple simultaneous trials, depending on the number of parallel workers
% available.
%
% A table of results displays the training loss and validation accuracy for
% each trial.
%
% <<../exp_mgr_gan_results.png>>
%
% While the experiment is running, click *Training Plot* to display the
% training plot and track the progress of each trial.
% 
% <<../exp_mgr_gan_training.png>>

%% Evaluate Results
% To find the best result for your experiment, sort the table of results
% using the combined score.
% 
% # Point to the *scoreCombined* column. 
% # Click the triangle icon.
% # Select *Sort in Descending Order*.
%
% The trial with the highest combined score appears at the top of the
% results table.
%
% <<../exp_mgr_gan_results_sorted.png>>
%
% Evaluate the quality of the GAN by generating and inspecting the images
% produced by the trained generator.
%
% # Select the trial with the highest combined score.
% # On the *Experiment Manager* toolstrip, click *Export*. 
% # In the dialog window, enter the name of a workspace variable for the
% exported training output. The default name is |trainingOutput|.
% # Test the trained generator network by calling the |generateTestImages|
% function, which is listed in <#9 Appendix 3> at the end of this example.
% Use the exported training output as the input to the function. For
% instance, in the MATLAB Command Window, enter:
%
%  generateTestImages(trainingOutput)
%
% The function creates a batch of 25 random vectors to input to the
% generator network and displays the resulting images.
%
% <<../exp_mgr_gan_test_images1.png>>
%
% Using the combined score to sort your results might not identify the best
% trial in all cases. For best results, repeat this process for each trial
% with a high combined score, visually checking that the generator produces
% a variety of images without many duplicates. If the images have little
% diversity and some of them are almost identical, then your generator is
% likely affected by mode collapse. For more information, see
% <docid:nnet_ug#mw_cedc8b1d-2b1e-44ef-bed8-988c8139016f Mode Collapse>.
%
% To record observations about the results of your experiment, add an
% annotation. 
%
% # In the results table, right-click the *scoreCombined* cell for
% the best trial.
% # Select *Add Annotation*. 
% # In the *Annotations* pane, enter your observations in the text box.
%
% <<../exp_mgr_gan_annotation.png>>
%
% For more information, see
% <docid:nnet_ref#mw_ffdb175d-cd27-4954-8fa6-e006dbb6dfd0 Sort, Filter, and
% Annotate Experiment Results>.

%% Rerun Experiment
% After you identify the combination of hyperparameters that generates the
% best images, run the experiment a second time to train the network for a
% longer period of time. 
%
% # Return to the experiment definition pane.
% # In the hyperparameter table, enter the hyperparameter values from your
% best trial. For example, to use the values from trial 3, change the value
% of |dropoutProb| to |0.75| and |flipFactor| to |0.1|.
% # Open the training function and specify a longer training time. Under
% *Specify Training Options*, change the value of |numEpochs| to |500|.
% # Run the experiment using the new hyperparameter values and training
% function. Experiment Manager runs a single trial. Training takes about 10
% times longer than the previous trials.
% # When the experiment finishes, export the training output and run the
% |generateTestImages| function to test the new generator network. As
% before, visually check that the generator produces a variety of images
% without many duplicates.
%
% <<../exp_mgr_gan_test_images2.png>>

%% Close Experiment
% In the *Experiment Browser* pane, right-click the name of the project and
% select *Close Project*. Experiment Manager closes all of the experiments
% and results contained in the project.

%% Appendix 1: Training Function
% This function specifies the training data, network architecture, training
% options, and training procedure used by the experiment.
%
% *Input*
%
% * |params| is a structure with fields from the Experiment Manager
% hyperparameter table.
% * |monitor| is an |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 generator network, the 
% trained discriminator network, and the execution environment used for
% training. Experiment Manager saves this output, so you can export it to
% the MATLAB workspace when the training is complete.
%
% <include>exp_mgr_gan_training_function.m</include>

%% Appendix 2: Custom Training Helper Functions
% 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

%% Appendix 3: Generate Test Images
% This function creates a batch of 25 random vectors to input to the
% generator network and displays the resulting images. Use this function to
% visually check that the generator produces a variety of images without
% many duplicates. If the images have little diversity and some of them are
% almost identical, then your generator is likely affected by mode
% collapse.
%
% <include>exp_mgr_gan_generateTestImages.m</include>