%% Train Conditional Generative Adversarial Network (CGAN)
% This example shows how to train a conditional generative adversarial network (CGAN) to generate images.
% A generative adversarial network (GAN) is a type of deep learning network that can generate data with similar characteristics as the input training data.
% A GAN consists of two networks that train together:
% Generator — Given a vector of random values as input, this network generates data with the same structure as the training data.
% Discriminator  — 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".
%
% A conditional generative adversarial network is a type of GAN that also takes advantage of labels during the training process.
% The generator - Given a label and random array as input, this network generates data with the same structure as the training data observations corresponding to the same label.
% The discriminator - Given batches of labeled 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".
%
% To train a conditional GAN, train both networks simultaneously to maximize the performance of both:
% Train the generator to generate data that "fools" the discriminator.
% Train the discriminator to distinguish between real and generated data.
% To maximize the performance of the generator, maximize the loss of the discriminator when given generated labeled data. That is, the objective of the generator is to generate labeled data that the discriminator classifies as "real".
% To maximize the performance of the discriminator, minimize the loss of the discriminator when given batches of both real and generated labeled data. That is, 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 that corresponds to the input labels and a discriminator that has learned strong feature representations that are characteristic of the training data for each label.
%% Load Training Data
% Download and extract the Flowers data set [1].

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

% Create an image datastore containing the photos.

datasetFolder = fullfile(imageFolder);

imds = imageDatastore(datasetFolder, ...
    'IncludeSubfolders',true, ...
    'LabelSource','foldernames');

% View the number of classes.

classes = categories(imds.Labels);
numClasses = numel(classes)

% Augment the data to include random horizontal flipping and resize the images to have size 64-by-64.

augmenter = imageDataAugmenter('RandXReflection',true);
augimds = augmentedImageDatastore([64 64],imds,'DataAugmentation',augmenter);

%
%% Define Generator Network
% Define the following two-input network, which generates images given 1-by-1-by-100 arrays of random values and corresponding labels:
%
% This network:
% Converts the 1-by-1-by-100 arrays of noise to 4-by-4-by-1024 arrays.
% Converts the categorical labels to embedding vectors and reshapes them to a 4-by-4 array.
% Concatenates the resulting images from the two inputs along the channel dimension. The output is a 4-by-4-by-1025 array.
% Upscales the resulting arrays to 64-by-64-by-3 arrays using a series of transposed convolution layers with batch normalization and ReLU layers.
% Define this network architecture as a layer graph and specify the following network properties.
% For the categorical inputs, use an embedding dimension of 50.
% For the transposed convolution layers, specify 5-by-5 filters with a decreasing number of filters for each layer, a stride of 2, and "same" cropping of the output.
% For the final transposed convolution layer, specify a three 5-by-5 filter corresponding to the three RGB channels of the generated images.
% At the end of the network, include a tanh layer.
% To project and reshape the noise input, use the custom layer projectAndReshapeLayer, attached to this example as a supporting file. The projectAndReshapeLayer object upscales the input using a fully connected operation and reshapes the output to the specified size.
% To input the labels into the network, use an imageInputLayer object and specify an image size of 1-by-1. To embed and reshape the label input, use the custom layer embedAndReshapeLayer, attached to this example as a supporting file. The embedAndReshapeLayer object converts a categorical label to a one-channel image of the specified size using an embedding and a fully connected operation.

numLatentInputs = 100;
embeddingDimension = 50;
numFilters = 64;

filterSize = 5;
projectionSize = [4 4 1024];

layersGenerator = [
    imageInputLayer([1 1 numLatentInputs],'Normalization','none','Name','noise')
    projectAndReshapeLayer(projectionSize,numLatentInputs,'proj');
    concatenationLayer(3,2,'Name','cat');
    transposedConv2dLayer(filterSize,4*numFilters,'Name','tconv1')
    batchNormalizationLayer('Name','bn1')
    reluLayer('Name','relu1')
    transposedConv2dLayer(filterSize,2*numFilters,'Stride',2,'Cropping','same','Name','tconv2')
    batchNormalizationLayer('Name','bn2')
    reluLayer('Name','relu2')
    transposedConv2dLayer(filterSize,numFilters,'Stride',2,'Cropping','same','Name','tconv3')
    batchNormalizationLayer('Name','bn3')
    reluLayer('Name','relu3')
    transposedConv2dLayer(filterSize,3,'Stride',2,'Cropping','same','Name','tconv4')
    tanhLayer('Name','tanh')];

lgraphGenerator = layerGraph(layersGenerator);

layers = [
    imageInputLayer([1 1],'Name','labels','Normalization','none')
    embedAndReshapeLayer(projectionSize(1:2),embeddingDimension,numClasses,'emb')];

lgraphGenerator = addLayers(lgraphGenerator,layers);
lgraphGenerator = connectLayers(lgraphGenerator,'emb','cat/in2');

% To train the network with a custom training loop and enable automatic differentiation, convert the layer graph to a dlnetwork object.

dlnetGenerator = dlnetwork(lgraphGenerator)

%
%% Define Discriminator Network
% Define the following two-input network, which classifies real and generated 64-by-64 images given a set of images and the corresponding labels.
%
% Create a network that takes as input 64-by-64-by-1 images and the corresponding labels and outputs a scalar prediction score using a series of convolution layers with batch normalization and leaky ReLU layers. Add noise to the input images using dropout.
% For the dropout layer, specify a dropout probability of 0.75.
% For the convolution layers, specify 5-by-5 filters with an increasing number of filters for each layer. Also specify a stride of 2 and to padding of the output on each edge.
% For the leaky ReLU layers, specify a scale of 0.2.
% For the final layer, specify a convolution layer with one 4-by-4 filter.

dropoutProb = 0.75;
numFilters = 64;
scale = 0.2;

inputSize = [64 64 3];
filterSize = 5;

layersDiscriminator = [
    imageInputLayer(inputSize,'Normalization','none','Name','images')
    dropoutLayer(dropoutProb,'Name','dropout')
    concatenationLayer(3,2,'Name','cat')
    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);

layers = [
    imageInputLayer([1 1],'Name','labels','Normalization','none')
    embedAndReshapeLayer(inputSize,embeddingDimension,numClasses,'emb')];

lgraphDiscriminator = addLayers(lgraphDiscriminator,layers);
lgraphDiscriminator = connectLayers(lgraphDiscriminator,'emb','cat/in2');

% To train the network with a custom training loop and enable automatic differentiation, convert the layer graph to a dlnetwork object.

dlnetDiscriminator = dlnetwork(lgraphDiscriminator)

%
%% Define Model Gradients and Loss Functions
% Create the function modelGradients, listed in the Model Gradients Function section of the example, which takes as input the generator and discriminator networks, a mini-batch of input data, and an array of random values, and returns the gradients of the loss with respect to the learnable parameters in the networks and an array of generated images.
%
%% Specify Training Options
% Train with a mini-batch size of 128 for 500 epochs.

numEpochs = 500;
miniBatchSize = 128;
augimds.MiniBatchSize = miniBatchSize;

% Specify the options for Adam optimization. For both networks, use:
% A learning rate of 0.0002
% A gradient decay factor of 0.5
% A squared gradient decay factor of 0.999

learnRate = 0.0002;
gradientDecayFactor = 0.5;
squaredGradientDecayFactor = 0.999;

% Train on a GPU if one is available. Using a GPU requires Parallel Computing Toolbox™ and a CUDA® enabled NVIDIA® GPU with compute capability 3.0 or higher.

executionEnvironment = "auto";

% Update the training progress plots every 100 iterations.

validationFrequency = 100;

% If the discriminator learns to discriminate between real and generated images too quickly, then the generator may fail to train. To better balance the learning of the discriminator and the generator, randomly flip the labels of a proportion of the real images. Specify a flip factor of 0.5.

flipFactor = 0.5;

%
%% Train Model
% Train the model using a custom training loop. Loop over the training data and update the network parameters at each iteration. To monitor the training progress, display a batch of generated images using a held-out array of random values to input into the generator and the network scores.
% Initialize the parameters for the Adam optimizer.

velocityDiscriminator = [];
trailingAvgGenerator = [];
trailingAvgSqGenerator = [];
trailingAvgDiscriminator = [];
trailingAvgSqDiscriminator = [];

% Initialize the plot of the training progress. Create a figure and resize it to have twice the width.

f = figure;
f.Position(3) = 2*f.Position(3);

% Create subplots of the generated images and of the scores plot.

imageAxes = subplot(1,2,1);
scoreAxes = subplot(1,2,2);

% Initialize animated lines for the scores plot.

lineScoreGenerator = animatedline(scoreAxes,'Color',[0 0.447 0.741]);
lineScoreDiscriminator = animatedline(scoreAxes, 'Color', [0.85 0.325 0.098]);

% Customize the appearance of the plots.

legend('Generator','Discriminator');
ylim([0 1])
xlabel("Iteration")
ylabel("Score")
grid on

% To monitor training progress, create a held-out batch of 25 1-by-1-by-numLatentInputs arrays of random values and a corresponding set of labels 1 through 5 (corresponding to the classes) repeated 5 times, where the labels are in the fourth dimension of the array.

numValidationImagesPerClass = 5;
ZValidation = randn(1,1,numLatentInputs,numValidationImagesPerClass*numClasses,'single');
TValidation = single(repmat(1:numClasses,[1 numValidationImagesPerClass]));
TValidation = permute(TValidation,[1 3 4 2]);

% Convert the data to dlarray objects and specify the dimension labels 'SSCB' (spatial, spatial, channel, batch).

dlZValidation = dlarray(ZValidation, 'SSCB');
dlTValidation = dlarray(TValidation, 'SSCB');

% For GPU training, convert the data to gpuArray objects.

if (executionEnvironment == "auto" && canUseGPU) || executionEnvironment == "gpu"
    dlZValidation = gpuArray(dlZValidation);
    dlTValidation = gpuArray(dlTValidation);
end

% Train the GAN. For each epoch, shuffle the data and loop over mini-batches of data.
% For each mini-batch:
% To ensure that the inputs to the discriminator match the outputs of the generator, rescale the real images so that the pixels take values in the range [-1, 1].
% Convert the image data and labels to dlarray objects with underlying type single and specify the dimension labels 'SSCB' (spatial, spatial, channel, batch).
% Generate a dlarray object containing an array of random values for the generator network.
% For GPU training, convert the data to gpuArray objects.
% Evaluate the model gradients using dlfeval and the modelGradients function.
% Update the network parameters using the adamupdate function.
% Plot the scores of the two networks.
% After every validationFrequency iterations, display a batch of generated images for a fixed held-out generator input.
% Training can take some time to run.

iteration = 0;
start = tic;

% Loop over epochs.
for epoch = 1:numEpochs
    
    % Reset and shuffle datastore.
    reset(augimds);
    augimds = shuffle(augimds);
    
    % Loop over mini-batches.
    while hasdata(augimds)
        iteration = iteration + 1;
        
        % Read mini-batch of data and generate latent inputs for the
        % generator network.
        data = read(augimds);
        
        % Ignore last partial mini-batch of epoch.
        if size(data,1) < miniBatchSize
            continue
        end
        
        X = cat(4,data{:,1}{:});
        X = single(X);
        
        T = single(data.response);
        T = permute(T,[2 3 4 1]);
        
        Z = randn(1,1,numLatentInputs,miniBatchSize,'single');
        
        % Rescale the images in the range [-1 1].
        X = rescale(X,-1,1,'InputMin',0,'InputMax',255);
        
        % Convert mini-batch of data to dlarray and specify the dimension labels
        % 'SSCB' (spatial, spatial, channel, batch).
        dlX = dlarray(X, 'SSCB');
        dlZ = dlarray(Z, 'SSCB');
        dlT = dlarray(T, 'SSCB');
        
        % If training on a GPU, then convert data to gpuArray.
        if (executionEnvironment == "auto" && canUseGPU) || executionEnvironment == "gpu"
            dlX = gpuArray(dlX);
            dlZ = gpuArray(dlZ);
            dlT = gpuArray(dlT);
        end
        
        % Evaluate the model gradients and the generator state using
        % dlfeval and the modelGradients function listed at the end of the
        % example.
        [gradientsGenerator, gradientsDiscriminator, stateGenerator, scoreGenerator, scoreDiscriminator] = ...
            dlfeval(@modelGradients, dlnetGenerator, dlnetDiscriminator, dlX, dlT, dlZ, flipFactor);
        dlnetGenerator.State = stateGenerator;
        
        % Update the discriminator network parameters.
        [dlnetDiscriminator,trailingAvgDiscriminator,trailingAvgSqDiscriminator] = ...
            adamupdate(dlnetDiscriminator, gradientsDiscriminator, ...
            trailingAvgDiscriminator, trailingAvgSqDiscriminator, iteration, ...
            learnRate, gradientDecayFactor, squaredGradientDecayFactor);
        
        % Update the generator network parameters.
        [dlnetGenerator,trailingAvgGenerator,trailingAvgSqGenerator] = ...
            adamupdate(dlnetGenerator, gradientsGenerator, ...
            trailingAvgGenerator, trailingAvgSqGenerator, iteration, ...
            learnRate, gradientDecayFactor, squaredGradientDecayFactor);
        
        % Every validationFrequency iterations, display batch of generated images using the
        % held-out generator input.
        if mod(iteration,validationFrequency) == 0 || iteration == 1
            
            % Generate images using the held-out generator input.
            dlXGeneratedValidation = predict(dlnetGenerator,dlZValidation,dlTValidation);
            
            % Tile and rescale the images in the range [0 1].
            I = imtile(extractdata(dlXGeneratedValidation), ...
                'GridSize',[numValidationImagesPerClass numClasses]);
            I = rescale(I);
            
            % Display the images.
            subplot(1,2,1);
            image(imageAxes,I)
            xticklabels([]);
            yticklabels([]);
            title("Generated Images");
        end
        
        % Update the scores plot
        subplot(1,2,2)
        addpoints(lineScoreGenerator,iteration,...
            double(gather(extractdata(scoreGenerator))));
        
        addpoints(lineScoreDiscriminator,iteration,...
            double(gather(extractdata(scoreDiscriminator))));
        
        % Update the title with training progress information.
        D = duration(0,0,toc(start),'Format','hh:mm:ss');
        title(...
            "Epoch: " + epoch + ", " + ...
            "Iteration: " + iteration + ", " + ...
            "Elapsed: " + string(D))
        
        drawnow
    end
end

% Here, the discriminator has learned a strong feature representation that identifies real images among generated images. In turn, the generator has learned a similarly strong feature representation that allows it to generate realistic looking data. Each column corresponds to a single class.
% The training plot shows the scores of the generator and discriminator networks. To learn more about how to interpret the network scores, see Monitor GAN Training Progress and Identify Common Failure Modes.
%
%% Generate New Images
% To generate new images of a particular class, use the predict function on the generator with a dlarray object containing a batch of 1-by-1-by-numLatentInputs arrays of random values and an array of labels corresponding to the desired classes. Convert the data to dlarray objects and specify the dimension labels 'SSCB' (spatial, spatial, channel, batch). For GPU prediction, convert the data to gpuArray. To display the images together, use the imtile function and rescale the images using the rescale function.
% Create an array of 36 vectors of random values corresponding to the first class.

numObservationsNew = 36;
idxClass = 1;
Z = randn(1,1,numLatentInputs,numObservationsNew,'single');
T = repmat(single(idxClass),[1 1 1 numObservationsNew]);

% Convert the data to dlarray objects with the dimension labels 'SSCB' (spatial, spatial, channels, batch).

dlZ = dlarray(Z,'SSCB');
dlT = dlarray(T,'SSCB');

%  To generate images using the GPU, also convert the data to gpuArray objects.

if (executionEnvironment == "auto" && canUseGPU) || executionEnvironment == "gpu"
    dlZ = gpuArray(dlZ);
    dlT = gpuArray(dlT);
end

% Generate images using the predict function with the generator network.

dlXGenerated = predict(dlnetGenerator,dlZ,dlT);

% Display the generated images in a plot.

figure
I = imtile(extractdata(dlXGenerated));
I = rescale(I);
imshow(I)
title("Class: " + classes(idxClass))

% Here, the generator network generates images conditioned on the specified class.
%
%% Model Gradients Function
% The function modelGradients takes as input the generator and discriminator dlnetwork objects dlnetGenerator and dlnetDiscriminator, a mini-batch of input data dlX, the corresponding labels dlT, and an array of random values dlZ, and returns the gradients of the loss with respect to the learnable parameters in the networks, the generator state, and the network scores.
% If the discriminator learns to discriminate between real and generated images too quickly, then the generator may fail to train. To better balance the learning of the discriminator and the generator, randomly flip the labels of a proportion of the real images.

function [gradientsGenerator, gradientsDiscriminator, stateGenerator, scoreGenerator, scoreDiscriminator] = ...
    modelGradients(dlnetGenerator, dlnetDiscriminator, dlX, dlT, dlZ, flipFactor)

% Calculate the predictions for real data with the discriminator network.
dlYPred = forward(dlnetDiscriminator, dlX, dlT);

% Calculate the predictions for generated data with the discriminator network.
[dlXGenerated,stateGenerator] = forward(dlnetGenerator, dlZ, dlT);
dlYPredGenerated = forward(dlnetDiscriminator, dlXGenerated, dlT);

% Calculate probabilities.
probGenerated = sigmoid(dlYPredGenerated);
probReal = sigmoid(dlYPred);

% Calculate the generator and discriminator scores
scoreGenerator = mean(probGenerated);
scoreDiscriminator = (mean(probReal) + mean(1-probGenerated)) / 2;

% Flip labels.
numObservations = size(dlYPred,4);
idx = randperm(numObservations,floor(flipFactor * numObservations));
probReal(:,:,:,idx) = 1 - probReal(:,:,:,idx);

% Calculate the GAN loss.
[lossGenerator, lossDiscriminator] = ganLoss(probReal, probGenerated);

% For each network, calculate the gradients with respect to the loss.
gradientsGenerator = dlgradient(lossGenerator, dlnetGenerator.Learnables,'RetainData',true);
gradientsDiscriminator = dlgradient(lossDiscriminator, dlnetDiscriminator.Learnables);

end

%% GAN Loss Function
% The objective of the generator is to generate data that the discriminator classifies as "real". To maximize the probability that images from the generator are classified as real by the discriminator, minimize the negative log likelihood function.
% Given the output Y of the discriminator:
% \hat{Y} =\sigma \left(Y\right) is the probability that the input image belongs to the class "real".
% 1-\hat{Y} is the probability that the input image belongs to the class "generated".
% Note that the sigmoid operation \sigma in the modelGradients function. The loss function for the generator is given by
%     \textrm{lossGenerator}=-\textrm{mean}\left(\log \left({\hat{Y} }_{\textrm{Generated}} \right)\right),
% where \hat{Y}_{Generated} contains the discriminator output probabilities for the generated images.
% The objective of the discriminator is to not be "fooled" by the generator. To maximize the probability that the discriminator successfully discriminates between the real and generated images, minimize the sum of the corresponding negative log likelihood functions. The loss function for the discriminator is given by
%     \textrm{lossDiscriminator}=-\textrm{mean}\left(\log \left({\hat{Y} }_{\textrm{Real}} \right)\right)-\textrm{mean}\left(\log \left(1-{\hat{Y} }_{\textrm{Generated}} \right)\right),
% where \hat{Y}_{Real} contains the discriminator output probabilities for the real images.

function [lossGenerator, lossDiscriminator] = ganLoss(scoresReal,scoresGenerated)

% Calculate losses for the discriminator network.
lossGenerated = -mean(log(1 - scoresGenerated));
lossReal = -mean(log(scoresReal));

% Combine the losses for the discriminator network.
lossDiscriminator = lossReal + lossGenerated;

% Calculate the loss for the generator network.
lossGenerator = -mean(log(scoresGenerated));

end

%% References
% The TensorFlow Team. Flowers http://download.tensorflow.org/example_images/flower_photos.tgz
% Copyright 2019 The MathWorks, Inc.