%% Parameter Pruning and Quantization of Image Classification Network
% This example shows how to prune the parameters of a trained neural network using two different parameter score metrics: magnitude score [1] and Synaptic Flow score [2].
% For many applications when transfer learning is used to retrain an image classification network for a new task or when a new network is trained from scratch, the optimal network architecture is not known and might lead to overparameterization. This can result in the network having redundant connections. Pruning, also known as sparsification, is a compression technique that aims to identify these redundant, unnecessary connections that can be removed without affecting the network accuracy. When used in combination with network quantization, the inference speed and memory footprint can be reduced, making it easier to deploy the model.
% This example shows how to:
% Perform post-training, iterative, unstructured pruning without the need for training data
% Evaluate the performance of two different pruning algorithms
% Investigate the layer-wise sparsity induced via pruning
% Evaluate the impact of pruning on classification accuracy
% Evaluate the impact of quantization on the accuracy of the pruned network
% This example uses a simple convolutional neural network to classify hand-written digits 0-9. For more information on setting up the data used for training and validation, see Create Simple Deep Learning Network for Classification.
%% Load Pretrained Network and Data
% Load the training and validation data. Train a convolutional neural network for the classification task.

[imdsTrain, imdsValidation] = loadDigitDataset();
net = trainDigitDataNetwork(imdsTrain, imdsValidation);
trueLabels = imdsValidation.Labels;
classes = categories(trueLabels);

% Create a minibatchqueue object containing the validation data. Specify whether to evalute the network on a GPU, if one is available. By default, the minibatchqueue object converts each output to a gpuArray if a GPU is available. Using a GPU requires Parallel Computing Toolbox™ and a supported GPU device. For information on supported devices, see GPU Support by Release (Parallel Computing Toolbox).

executionEnvironment = "auto";
miniBatchSize = 128;
imdsValidation.ReadSize = miniBatchSize;
mbqValidation = minibatchqueue(imdsValidation,1,...
    'MiniBatchSize',miniBatchSize,...
    'MiniBatchFormat','SSCB',...
    'MiniBatchFcn',@preprocessMiniBatch,...
    'OutputEnvironment',executionEnvironment);

%%
%
%% Neural Network Pruning
% The goal of neural network pruning is to identify and remove dead neurons or unimportant connections in order to reduce the size of the network without affecting network accuracy. In the figure below, the original network has connections mapping each neuron to the next layer neurons. The pruned network has many fewer connections than the original network.
%
% A pruning algorithm assigns score to each parameter in the network which corresponds to the importance of a specific connection or neuron in the network. There are two pruning approaches to achieve a target percent sparsity:
% One shot pruning - Remove a specified percentage of connections based on their score in one step. This method is prone to layer collapse when reaching high sparsity.
% Iterative pruning - Acheive the target sparisty in a series of iterative steps. This method is beneficial when evaluated scores are sensitive to network structure. As scores are reevaluated at every iteration, the parameter ranking will take into account the algorithm sensitivity to network structure as connections are removed at every iteration.
% In this example, the iterative pruning method is used to achieve a target sparsity.
%% Iterative Pruning
%
%
%% Convert to dlnetwork
% Some pruning algorithms require the construction of custom cost functions and evaluating gradients with respect to the cost function to calculate the paramter score. To create a custom cost function, first convert the pretrained network to a dlnetwork.
% Convert the network to a layer graph and remove the layers used for classification using removeLayers. For the current example, the probability or softmax layer is removed for the synaptic flow pruning implementation.

lgraph = layerGraph(net.Layers);
lgraph = removeLayers(lgraph,["softmax","classoutput"]);
dlnet = dlnetwork(lgraph);

% You can use analyzeNetwork to analyze the network architecture and learnables.

analyzeNetwork(dlnet)

% Evaluate the accuracy of the network before pruning.

accuracyOriginalNet = evaluateAccuracy(dlnet,mbqValidation,classes,trueLabels)

% The layers with learnable parameters are the 3 convolution layers and one fully connected layer. The network initially consists of total 21578 learnable paramters.

numTotalParams = sum(cellfun(@numel,dlnet.Learnables.Value))
numNonZeroPerParam = cellfun(@(w)nnz(extractdata(w)),dlnet.Learnables.Value)

% Sparsity is defined as the percentage of parameters in the network with a value of zero.

initialSparsity = 1-sum(numNonZeroPerParam)/numTotalParams

% Before pruning, the network has a sparisty of zero.
%% Create Iteration Scheme
% To define an iterative pruning scheme, specify the target sparsity and number of iterations. This example uses linearly-spaced iterations to achieve the target sparsity.

numIterations = 10; 
targetSparsity = 0.90;
iterationScheme = linspace(0,targetSparsity,numIterations); 

%% Pruning Loop
% For each iteration, the custom pruning loop in this example performs the following steps:
% Apply pruning mask to learnable parameters of the network
% Calculate the score for each connection
% Rank the scores for all connections in the network based on the pruning algorithm of choice
% Determine the threshold for removing connections with lowest scores
% Create pruning mask using the threshold
%% Network Mask
% Instead of setting entries in the weight arrays directly to zero, the pruning algorithm creates a binary mask for each learnable parameter that specifies whether a connection is pruned. The mask allows you to explore the behavior of the pruned network and try different pruining schemes without changing the underlying network structure.
% For example, consider the following weights.

testWeight = [10.4 5.6 0.8 9];

% Create a binary mask for each parameter in |testWeight|.

testMask = [1 0 1 0];

% Apply the mask to |testWeight| to get the pruned weights.

testWeightsPruned = testWeight.*testMask

% For iterative pruning, a binary mask is created for each iteration that contains pruning information for each iteration step. Applying the mask to the weights array does not change the size of the array nor the structure of the neural network. Therefore, the pruning step does not directly result in any speedup during inference or compression of the network size on disk.
% Initialize a plot that compares the performance of the pruned network to the original network.

figure
plot(100*iterationScheme([1,end]),100*accuracyOriginalNet*[1 1],'*-b','LineWidth',2,"Color","b")
ylim([0 100])
xlim(100*iterationScheme([1,end]))
xlabel("Sparsity (%)")
ylabel("Accuracy (%)")
legend("Original Accuracy","Location","southwest")
grid on

%
%% Magnitude Pruning
% Magnitude pruning [1] assigns a score to each parameter equal to its absolute value. It is assumed that the absolute value of a parameter corresponds to its relative importance to the accuracy of the network.
% Initialize the mask. For the first iteration, no parameters are pruned and the sparsity is 0%.

pruningMaskMagnitude = cell(1,numIterations); 
pruningMaskMagnitude{1} = dlupdate(@(p)true(size(p)), dlnet.Learnables);

% Below is an implementation of magnitude pruning. Iterative pruning allows the flexibility to trace back a pruned network based on its accuracy.

lineAccuracyPruningMagnitude = animatedline('Color','g','Marker','o','LineWidth',1.5);
legend("Original Accuracy","Magnitude Pruning Accuracy","Location","southwest")
    
prunedNetMagnitude = dlnet;
scoresMagnitude = calculateMagnitudeScore(prunedNetMagnitude);

for idx = 1:numel(iterationScheme)
    
    % Update the pruning mask
    pruningMaskMagnitude{idx} = calculateMask(scoresMagnitude,iterationScheme(idx));
    
    % Check the number of zero entries in the pruning mask
    numPrunedParams = sum(cellfun(@(m)nnz(~extractdata(m)),pruningMaskMagnitude{idx}.Value));
    sparsity = numPrunedParams/numTotalParams;
    
    % Apply pruning mask to network parameters
    prunedNetMagnitude.Learnables = dlupdate(@(W,M)W.*M, prunedNetMagnitude.Learnables, pruningMaskMagnitude{idx});
    
    % Compute validation accuracy on pruned network
    accuracyMagnitude = evaluateAccuracy(prunedNetMagnitude,mbqValidation,classes,trueLabels);
    
    % Display the pruning progress
    addpoints(lineAccuracyPruningMagnitude,100*sparsity,100*accuracyMagnitude)
    title("Pruning Performance")
    drawnow
end

%
%% SynFlow Pruning
% Synaptic flow conservation (SynFlow) [2] scores are used for pruning. This can be applied to networks that use linear activation functions like ReLU.
% Initialize the mask. For the first iteration, no parameters are pruned and the sparsity is 0%.

pruningMaskSynFlow = cell(1,numIterations); 
pruningMaskSynFlow{1} = dlupdate(@(p)true(size(p)), dlnet.Learnables);

% The input data used to compute the scores is a single image containing ones.

dlX = dlarray(ones(net.Layers(1).InputSize),'SSC');

% If you are using a GPU, convert the data to a gpuArray.

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

% The below loop implements iterative synaptic flow score for pruning[2] where a custom cost function evaluates the synflow score for each parameter used for network pruning

lineAccuracyPruningSynflow = animatedline('Color','r','Marker','o','LineWidth',1.5);
legend("Original Accuracy","Magnitude Pruning Accuracy","Synaptic Flow Accuracy","Location","southwest")

prunedNetSynFlow = dlnet;

% Iteratively increase sparsity
for idx = 1:numel(iterationScheme)
    % Compute SynFlow scores
    scoresSynFlow = calculateSynFlowScore(prunedNetSynFlow,dlX);
    
    % Update the pruning mask
    pruningMaskSynFlow{idx} = calculateMask(scoresSynFlow,iterationScheme(idx));
    
    % Check the number of zero entries in the pruning mask
    numPrunedParams = sum(cellfun(@(m)nnz(~extractdata(m)),pruningMaskSynFlow{idx}.Value));
    sparsity = numPrunedParams/numTotalParams;
    
    % Apply pruning mask to network parameters
    prunedNetSynFlow.Learnables = dlupdate(@(W,M)W.*M, prunedNetSynFlow.Learnables, pruningMaskSynFlow{idx});
    
    % Compute validation accuracy on pruned network
    accuracySynFlow = evaluateAccuracy(prunedNetSynFlow,mbqValidation,classes,trueLabels);
     
    % Display the pruning progress
    addpoints(lineAccuracyPruningSynflow,100*sparsity,100*accuracySynFlow)
    title("Pruning Performance")
    drawnow
end

%
%% Investigate Structure of Pruned Network
% Use the Sparsity vs Accuracy plot to determine which algorithm and sparsity level results in the least impact on the final accuracy metric.

pruningMethod = "SynFlow";
selectedIteration = 8;

prunedDLNet = createPrunedNet(dlnet,selectedIteration,pruningMaskSynFlow,pruningMaskMagnitude,pruningMethod);

[sparsityPerLayer, prunedChannelsPerLayer, numOutChannelsPerLayer, layerNames] = pruningStatistics(prunedDLNet);

% Earlier convolution layers are typically pruned less since they contain more relevant information about the core low-level structure of the image (e.g. edges and corners) which are essential for interpreting the image.
% Plot the sparisity per layer for the selected pruning method and iteration.

figure
bar(sparsityPerLayer*100)
title("Sparsity per layer")
xlabel("Layer")
ylabel("Sparsity (%)")
xticks(1:numel(sparsityPerLayer))
xticklabels(layerNames)
xtickangle(45)
set(gca,'TickLabelInterpreter','none')

% Even though the pruning algorithm can prune single connections, when the sparsity increases, whole filters or neurons are pruned in convolutional or fully connected layers. Removing complete filters or neurons can lead to significant strucutral size reduction.

figure
bar([prunedChannelsPerLayer,numOutChannelsPerLayer-prunedChannelsPerLayer],"stacked")
xlabel("Layer")
ylabel("Number of filters")
title("Number of filters per layer")
xticks(1:(numel(layerNames)))
xticklabels(layerNames)
xtickangle(45)
legend("Pruned number of channels/neurons" , "Original number of channels/neurons","Location","southoutside")
set(gca,'TickLabelInterpreter','none')

%
%% Evaluate Network Accuracy
% Compare the accuracy of the network before and after pruning.

YPredOriginal = modelPredictions(dlnet,mbqValidation,classes);
accOriginal = mean(YPredOriginal == trueLabels)
YPredPruned = modelPredictions(prunedDLNet,mbqValidation,classes);
accPruned = mean(YPredPruned == trueLabels)

% Create a confusion matrix chart to explore the true class labels to the predicted class labels for the original and pruned network.

figure
confusionchart(trueLabels,YPredOriginal, ...
    'ColumnSummary','column-normalized', ...
    'RowSummary','row-normalized');
title("Original Network")

%

figure("Position",[100 100 500 500])
confusionchart(trueLabels,YPredPruned, ...
    'ColumnSummary','column-normalized', ...
    'RowSummary','row-normalized');
title("Pruned Network")

% Pruning has a different impact on each label. Increasing pruning sparsity does not have the same impact on each label and might lead to significant degradation for specific labels.
%
%% Quantize Pruned Network
% Most networks trained in MATLAB use single precision floating point datatype which even for small networks can use significant memory and compute. Quantization as a compression technique converts the single precision datatypes to a lower precision datatype with minimal impact final network output. Deep Learning Model Quantization Library implements the quantization workflow in MATLAB. For a detailed example, see Quantize Residual Network Trained for Image Classification
% Pruning of the network impacts the range statistics of parameters and activations at each layer, which can lead to changed accuracy of the quantized network. To explore this difference, quantize the pruned network and use the quantized network to perform inference.
% Split the data into calibration and validation data sets.

calibrationDataStore = splitEachLabel(imdsTrain,0.1,'randomize');
validationDataStore = imdsValidation;

% Create a dlquantizer object and specify the pruned network as the network to quantize. Set "Simulation" to "on" to validate the quantized network by simulating in MATLAB.

prunedNet  = assembleNetwork([prunedDLNet.Layers ; net.Layers(end-1:end)]);

quantObjPrunedNetwork = dlquantizer(prunedNet,'ExecutionEnvironment','GPU'); 

% Use the calibrate function to exercise the network with the calibration data and collect range statistics for the weights, biases, and activations at each layer.

calResults = calibrate(quantObjPrunedNetwork, calibrationDataStore)

% Use the validate function to compare the results of the network before and after quantization using the validation data set.

valResults = validate(quantObjPrunedNetwork, validationDataStore);

% Examine the MetricResults.Result field of the validation output to see the accuracy of the quantized network.

valResults.MetricResults.Result
valResults.Statistics

%% Conclusion
% Pruning can be used to identify redundant connections in a network. This example demonstrates how to explore the behavior of a pruned network using masks, which do not alter the underlying structure of the network. Therefore, unless used in combination with sparsity-optimized libraries in deployment, this implementation of the floating-point pruned network using masks will not offer any specific compression or speedup in the final deployed solution. By using quantization to convert the floating-point network to an int8 implementation, a better deployment solution can be achieved.
%
%% Supporting Functions

function [X] = preprocessMiniBatch(XCell)
% Extract image data from cell and concatenate.
X = cat(4,XCell{:});
end

% Evaluate the accuracy of dlnetwork. Accuracy here is the percentage of labels correctly classified by the network.

function [accuracy] = evaluateAccuracy(dlnet,mbqValidation,classes,trueLabels)
YPred = modelPredictions(dlnet,mbqValidation,classes);
accuracy = mean(YPred == trueLabels);
end

% The calculateSynFlowScore function calculates SynFlow scores. The score is a product of the gradient
% \textrm{cost}=\sum f\left(\textrm{absolute}\left(\theta \right),\textrm{inputArray}\right);
% f is the function represented by the neural network
% \theta are the paramters of the network
% cost is the summation of all the values output from the last layer. This cost is used to evaluate the gradients on the paramters.
% \textrm{synFlowScore}=\frac{d\left(\textrm{cost}\right)*\theta }{d\theta }

function score = calculateSynFlowScore(dlnet,dlX)
dlnet.Learnables = dlupdate(@abs, dlnet.Learnables);
gradients = dlfeval(@modelGradients,dlnet,dlX);
score = dlupdate(@(g,w)g.*w, gradients, dlnet.Learnables);
end

function gradients = modelGradients(dlNet,inputArray)
% Evaluate the gradients on a given input to the dlnetwork
dlYPred = predict(dlNet,inputArray);
pseudoloss = sum(dlYPred,'all');
gradients = dlgradient(pseudoloss,dlNet.Learnables);
end

% The calculateMagnitudeScore function returns the magnitude score, defined as the element-wise absolute value of the parameters.

function score = calculateMagnitudeScore(dlnet)
score = dlupdate(@abs, dlnet.Learnables);
end

% The calculateMask function returns a binary mask for the network parameters based on the given scores and the target sparsity.

function[mask] = calculateMask(scoresMagnitude,sparsity)
% Evaluate the final mask based on parameter-wise scores and the final desired sparsity in the network.

% Create a score vector for ranking.
flattenedScores = cell2mat(cellfun(@(S)extractdata(gather(S(:))),scoresMagnitude.Value,'UniformOutput',false));
% Rank the scores and determine threshold for removing connections for the
% given sparsity
flattenedScores = sort(flattenedScores);
k = round(sparsity*numel(flattenedScores));
if k==0
    thresh = 0;
else
    thresh = flattenedScores(k);
end
% Create a binary mask 
mask = dlupdate( @(S)S>thresh, scoresMagnitude);
end

% The modelPredictions function for the given validation data evaluates the final classification output based on the predict call on the dlnetwork.

function predictions = modelPredictions(dlnet,mbq,classes)
predictions = [];
while hasdata(mbq)
    dlXTest = next(mbq);
    dlYPred = softmax(predict(dlnet,dlXTest));
    YPred = onehotdecode(dlYPred,classes,1)';
    predictions = [predictions; YPred];
end
reset(mbq)
end

% The createPrunedNet function returns the pruned dlnetwork for the specified pruning algorithm and iteration.

function prunedNet = createPrunedNet(dlnet,selectedIteration,pruningMaskSynFlow,pruningMaskMagnitude,pruningMethod)
switch pruningMethod
    case "Magnitude"
        prunedNet = dlupdate(@(W,M)W.*M, dlnet, pruningMaskMagnitude{selectedIteration});
    case "SynFlow"
        prunedNet = dlupdate(@(W,M)W.*M, dlnet, pruningMaskSynFlow{selectedIteration});
end
end

% The pruningStatistics function extracts detailed layer-level pruning statistics such as layer-level sparsity and the number of filters or neurons being pruned
% sparsityPerLayer - percentage of paramters pruned in each layer
% prunedChannelsPerLayer - number of channels/neurons in each layer that can be removed as a result of pruning
% numOutChannelsPerLayer - number of channels/neurons in each layer

function [sparsityPerLayer, prunedChannelsPerLayer, numOutChannelsPerLayer, layerNames] = pruningStatistics(dlnet)

layerNames = unique(dlnet.Learnables.Layer,'stable');
numLayers = numel(layerNames);
layerIDs = zeros(numLayers,1);
for idx = 1:numel(layerNames)
    layerIDs(idx) = find(layerNames(idx)=={dlnet.Layers.Name});
end

sparsityPerLayer = zeros(numLayers,1);
prunedChannelsPerLayer = zeros(numLayers,1);
numOutChannelsPerLayer = zeros(numLayers,1);

numParams = zeros(numLayers,1);
numPrunedParams = zeros(numLayers,1);
for idx = 1:numLayers
    layer = dlnet.Layers(layerIDs(idx));
    
    % Calculate the sparsity
    paramIDs = strcmp(dlnet.Learnables.Layer,layerNames(idx));
    paramValue = dlnet.Learnables.Value(paramIDs);
    for p = 1:numel(paramValue)
        numParams(idx) = numParams(idx) + numel(paramValue{p});
        numPrunedParams(idx) = numPrunedParams(idx) + nnz(extractdata(paramValue{p})==0);
    end

    % Calculate channel statistics
    sparsityPerLayer(idx) = numPrunedParams(idx)/numParams(idx);
    switch class(layer)
        case "nnet.cnn.layer.FullyConnectedLayer"
            numOutChannelsPerLayer(idx) = layer.OutputSize;
            prunedChannelsPerLayer(idx) = nnz(all(layer.Weights==0,2)&layer.Bias(:)==0);
        case "nnet.cnn.layer.Convolution2DLayer"
            numOutChannelsPerLayer(idx) = layer.NumFilters;
            prunedChannelsPerLayer(idx) = nnz(reshape(all(layer.Weights==0,[1,2,3]),[],1)&layer.Bias(:)==0);
        case "nnet.cnn.layer.GroupedConvolution2DLayer"
            numOutChannelsPerLayer(idx) = layer.NumGroups*layer.NumFiltersPerGroup;
            prunedChannelsPerLayer(idx) = nnz(reshape(all(layer.Weights==0,[1,2,3]),[],1)&layer.Bias(:)==0);
        otherwise
            error("Unknown layer: "+class(layer))
    end
end
end

% The loadDigitDatasetData function loads the digit data set into training and validation data.

function [imdsTrain, imdsValidation] = loadDigitDataset()
digitDatasetPath = fullfile(matlabroot,'toolbox','nnet','nndemos', ...
    'nndatasets','DigitDataset');
imds = imageDatastore(digitDatasetPath, ...
    'IncludeSubfolders',true,'LabelSource','foldernames');
[imdsTrain, imdsValidation] = splitEachLabel(imds,0.75,"randomized");
end

% The trainDigitDataNetwork trains a convolutional neural network on the training data.

function [net] = trainDigitDataNetwork(imdsTrain,imdsValidation)
layers = [
    imageInputLayer([28 28 1],"Normalization","rescale-zero-one")
    convolution2dLayer(3,8,'Padding','same')
    reluLayer
    maxPooling2dLayer(2,'Stride',2)
    
    convolution2dLayer(3,16,'Padding','same')
    reluLayer
    maxPooling2dLayer(2,'Stride',2)
    
    convolution2dLayer(3,32,'Padding','same')
    reluLayer
    
    fullyConnectedLayer(10)
    softmaxLayer
    classificationLayer];

% Specify the training options
options = trainingOptions('sgdm', ...
    'InitialLearnRate',0.01, ...
    'MaxEpochs',10, ...
    'Shuffle','every-epoch', ...
    'ValidationData',imdsValidation, ...
    'ValidationFrequency',30, ...
    'Verbose',false, ...
    'Plots','none',"ExecutionEnvironment","auto");

% Train network
net = trainNetwork(imdsTrain,layers,options);
end

%% References
% [1] Song Han, Jeff Pool, John Tran, and William J. Dally. "Learning both weights and connections for efficient neural networks." In Advances in Neural Information Processing Systems 28 (NIPS 2015), edited by C. Cortes, N. Lawrence, D. Lee, M. Sugiyama, and R. Garnett, 1135-1143.
% [2] Hidenori Tanaka, Daniel Kunin, Daniel L. K. Yamins, and Surya Ganguli. "Pruning neural networks without any data by iteratively conserving synaptic flow." 34th Conference on Neural Information Processing Systems (NeurlPS 2020)
% Copyright 2021 The MathWorks, Inc.