%% Accelerate Audio Deep Learning Using GPU-Based Feature Extraction
% In this example, you leverage GPUs for feature extraction and augmentation to decrease the time required to train a deep learning model. The model you train is a convolutional neural network (CNN) for acoustic fault recognition.
% Audio Toolbox™ includes gpuArray support for most feature extractors, including popular ones such as melSpectrogram and mfcc. For an overview of GPU support, see Code Generation and GPU Support.
%
%% Load Training Data
% Download and unzip the air compressor data set [1]. This data set consists of recordings from air compressors in a healthy state or one of seven faulty states.

url = 'https://www.mathworks.com/supportfiles/audio/AirCompressorDataset/AirCompressorDataset.zip';
downloadFolder = fullfile(tempdir,'AirCompressorDataset');
datasetLocation = tempdir;

if ~isfolder(fullfile(datasetLocation,'AirCompressorDataset'))
    loc = websave(downloadFolder,url);
    unzip(loc,fullfile(datasetLocation,'AirCompressorDataset'))
end

%
% Create an audioDatastore object to manage the data and split it into training and validation sets.

ads = audioDatastore(downloadFolder,'IncludeSubfolders',true,'LabelSource','foldernames');
rng default
[adsTrain,adsValidation] = splitEachLabel(ads,0.8);

%
% Visualize the number of files in the training and validation sets.

uniqueLabels = unique(adsTrain.Labels);
tblTrain = countEachLabel(adsTrain);
tblValidation = countEachLabel(adsValidation);
H = bar(uniqueLabels,[tblTrain.Count, tblValidation.Count],'stacked');
legend(H,["Training Set","Validation Set"],'Location','NorthEastOutside')

%
% Select random examples from the training set for plotting. Each recording has 50,000 samples sampled at 16 kHz.

t = (0:5e4-1)/16e3;
tiledlayout(4,2,'TileSpacing','compact','Padding','compact')
for n = 1:numel(uniqueLabels)
    idx = find(adsTrain.Labels==uniqueLabels(n));
    [x,fs] = audioread(adsTrain.Files{idx(randperm(numel(idx),1))});

    nexttile
    plotHandle = plot(t,x);
    if n == 7 || n == 8
        xlabel('Seconds');
    else
        set(gca,'xtick',[])
    end
    title(string(uniqueLabels(n)));
end

%
%% Preprocess Data on CPU and GPU
% In this example, you perform feature extraction and data augmentation while training the network. In this section, you define the feature extraction and augmentation pipeline and compare the speed of the pipeline executed on a CPU against the speed of the pipeline executed on a GPU. The output of this pipeline is the input to the CNN you train.
%
% Create an audioFeatureExtractor object to extract mel spectrums using 200 ms mel windows with a 5 ms hop. The output from extract is a numHops-by-128-by-1 array.

afe = audioFeatureExtractor('SampleRate',fs, ...
    'FFTLength',4096, ...
    'Window',hann(round(fs*0.2),'periodic'), ...
    'OverlapLength',round(fs*0.195), ...
    'melSpectrum',true);
setExtractorParams(afe,"melSpectrum",'NumBands',128);

featureVector = extract(afe,x);
[numHops,numFeatures,numChannels] = size(featureVector)

%
% Deep learning methods are data-hungry, and the training dataset in this example is relatively small. Use the mixup [2] augmentation technique to effectively enlarge the training set. In mixup, you merge the features extracted from two audio signals as a weighted sum. The two signals have different labels, and the label assigned to the merged feature matrix is probabilistically assigned based on the mixing coefficient. The mixup augmentation is implemented in the supporting object, Mixup.
% Create the pipeline to perform the following steps:
% Extract the log-mel spectrogram.
% Apply mixup to the feature matrices. The Mixup supporting object outputs a cell array containing the features and the label.
% Create two versions of the pipeline for comparison: one that executes the pipeline on your CPU, and one that converts the raw audio signal to a gpuArray so that the pipeline is executed on your GPU.

offset = eps;

adsTrainCPU = transform(adsTrain,@(x)log10(extract(afe,x)+offset));
mixerCPU = Mixup(adsTrainCPU);
adsTrainCPU = transform(adsTrainCPU,@(x,info)mix(mixerCPU,x,info),'IncludeInfo',true);

adsTrainGPU = transform(adsTrain,@gpuArray);
adsTrainGPU = transform(adsTrainGPU,@(x)log10(extract(afe,x)+offset));
mixerGPU = Mixup(adsTrainGPU);
adsTrainGPU = transform(adsTrainGPU,@(x,info)mix(mixerGPU,x,info),'IncludeInfo',true);

%
% For the validation set, apply the feature extraction pipeline but not the augmentation. Because you are not applying mixup, create a combined datastore to output a cell array containing the features and the label. Again, create one validation pipeline that executes on your GPU and one validation pipeline that executes on your CPU.

adsValidationGPU = transform(adsValidation,@gpuArray);
adsValidationGPU = transform(adsValidationGPU,@(x){log10(extract(afe,x)+offset)});
adsValidationGPU = combine(adsValidationGPU,arrayDatastore(adsValidation.Labels));

adsValidationCPU = transform(adsValidation,@(x){log10(extract(afe,x)+offset)});
adsValidationCPU = combine(adsValidationCPU,arrayDatastore(adsValidation.Labels));

%
% Compare the time it takes for the CPU and a single GPU to extract features and perform data augmentation.

tic
for ii = 1:numel(adsTrain.Files)
    x = read(adsTrainCPU);
end
cpuPipeline = toc;
reset(adsTrainCPU)

tic
for ii = 1:numel(adsTrain.Files)
    x = read(adsTrainGPU);
end
wait(gpuDevice) % Ensure all calculations are completed
gpuPipeline = toc;
reset(adsTrainGPU)

fprintf(['Read, extract, and augment train set (CPU): %0.2f seconds\n' ...
    'Read, extract, and augment train set (GPU): %0.2f seconds\n' ...
    'Speedup (CPU time)/(GPU time): %0.3f\n\n'], ...
    cpuPipeline,gpuPipeline,cpuPipeline/gpuPipeline)

%
% Reading from the datastore contributes a significant amount of the overall time to the pipeline. A comparison of just extraction and augmentation shows an even greater speedup. Compare just feature extraction on the GPU versus on the CPU.

x = read(ads);

extract(afe,x); % Incur initialization cost outside timing loop
tic
for ii = 1:numel(adsTrain.Files)
    features = log10(extract(afe,x)+offset);
end
cpuFeatureExtraction = toc;

x = gpuArray(x); % Incur initialization cost outside timing loop
extract(afe,x);
tic
for ii = 1:numel(adsTrain.Files)
    features = log10(extract(afe,x)+offset);
end
wait(gpuDevice) % Ensure all calculations are completed
gpuFeatureExtraction = toc;

fprintf(['Extract features from train set (CPU): %0.2f seconds\n' ...
    'Extract features from train set (GPU): %0.2f seconds\n' ...
    'Speedup (CPU time)/(GPU time): %0.3f\n\n'], ...
    cpuFeatureExtraction,gpuFeatureExtraction,cpuFeatureExtraction/gpuFeatureExtraction)

%
%% Define Network
% Define a convolutional neural network that takes the augmented mel spectrogram as input. This network applies a single convolutional layer consisting of 48 filters with 3-by-3 kernels, followed by a batch normalization layer and a ReLU activation layer. The time dimension is then collapsed using a max pooling layer. Finally, the output of the pooling layer is reduced using a fully connected layer followed by softmax and classification layers.  See List of Deep Learning Layers for more information.

numClasses = numel(categories(adsTrain.Labels));
imageSize = [numHops,afe.FeatureVectorLength];

layers = [
    imageInputLayer(imageSize,'Normalization','none')

    convolution2dLayer(3,48,'Padding','same')
    batchNormalizationLayer
    reluLayer

    maxPooling2dLayer([numHops,1])

    fullyConnectedLayer(numClasses)
    softmaxLayer
    classificationLayer('Classes',categories(adsTrain.Labels));
    ];

%
% To define the training options, use trainingOptions. Set the ExecutionEnvironment to multi-gpu to leverage multiple GPUs, if available. Otherwise, you can set ExecutionEnvironment to gpu. The computer used in this example has access to four Titan V GPU devices. In this example, the network training always leverages GPUs.

miniBatchSize = 128;
options = trainingOptions('adam', ...
    'Shuffle','every-epoch', ...
    'MaxEpochs',40, ...
    'LearnRateSchedule',"piecewise", ...
    'LearnRateDropPeriod',15, ...
    'LearnRateDropFactor',0.2, ...
    'MiniBatchSize',miniBatchSize, ...
    'Plots','training-progress', ...
    'Verbose',false, ...
    'ValidationData',adsValidationCPU, ...
    'ValidationFrequency',ceil(numel(adsTrain.Files)/miniBatchSize), ...
    'ExecutionEnvironment','multi-gpu');

%
%% Train Network
%% Train Network Using CPU-Based Preprocessing
% Call trainNetwork to train the network using your CPU for the feature extraction pipeline. The execution environment for the network training is your GPU(s).

tic
net = trainNetwork(adsTrainCPU,layers,options);
cpuTrainTime = toc;

%
%% Train Network Using GPU-Based Preprocessing
% Replace the validation data in the training options with the GPU-based pipeline. Train the network using your GPU(s) for the feature extraction pipeline. The execution environment for the network training is your GPU(s).

options.ValidationData = adsValidationGPU;
tic
net = trainNetwork(adsTrainGPU,layers,options);
gpuTrainTime = toc;

%
%% Compare CPU- and GPU-based Preprocessing
% Print the timing results for training using a CPU for feature extraction and augmentation, and training using GPU(s) for feature extraction and augmentation.

fprintf(['Training time (CPU): %0.2f seconds\n' ...
    'Training time (GPU): %0.2f seconds\n' ...
    'Speedup (CPU time)/(GPU time): %0.3f\n\n'], ...
    cpuTrainTime,gpuTrainTime,cpuTrainTime/gpuTrainTime)

%
%% Compare CPU and GPU Inference Performance
% Compare the time it takes to perform prediction on a single 3-second clip when feature extraction is performed on the GPU versus the CPU. In both cases, the network prediction happens on your GPU.

signalToClassify = read(ads);

gpuFeatureExtraction = gputimeit(@()predict(net,log10(extract(afe,gpuArray(signalToClassify))+offset)));
cpuFeatureExtraction = gputimeit(@()predict(net,log10(extract(afe,(signalToClassify))+offset)));

fprintf(['Prediction time for 3 s of data (feature extraction on CPU): %0.2f ms\n' ...
    'Prediction time for 3 s of data (feature extraction on GPU): %0.2f ms\n' ...
    'Speedup (CPU time)/(GPU time): %0.3f\n\n'], ...
    cpuFeatureExtraction*1e3,gpuFeatureExtraction*1e3,cpuFeatureExtraction/gpuFeatureExtraction)

%
% Compare the time it takes to perform prediction on a set of 3-second clips when feature extraction is performed on the GPU(s) versus the CPU. In both cases, the network prediction happens on your GPU(s).

adsValidationGPU = transform(adsValidation,@(x)gpuArray(x));
adsValidationGPU = transform(adsValidationGPU,@(x){log10(extract(afe,x)+offset)});
adsValidationCPU = transform(adsValidation,@(x){log10(extract(afe,x)+offset)});

gpuFeatureExtraction = gputimeit(@()predict(net,adsValidationGPU,'ExecutionEnvironment','multi-gpu'));
cpuFeatureExtraction = gputimeit(@()predict(net,adsValidationCPU,'ExecutionEnvironment','multi-gpu'));

fprintf(['Prediction time for validation set (feature extraction on CPU): %0.2f s\n' ...
    'Prediction time for validation set (feature extraction on GPU): %0.2f s\n' ...
    'Speedup (CPU time)/(GPU time): %0.3f\n\n'], ...
    cpuFeatureExtraction,gpuFeatureExtraction,cpuFeatureExtraction/gpuFeatureExtraction)

%
%% Conclusion
% It is well known that you can decrease the time it takes to train a network by leveraging GPU devices. This enables you to more quickly iterate and develop your final system. In many training setups, you can achieve additional performance gains by leveraging GPU devices for feature extraction and data augmentation. This example shows a significant decrease in the overall time it takes to train a CNN when leveraging GPU devices for feature extraction and data augmentation. Additionally, leveraging GPU devices for feature extraction at inference time, for both single-observations and data sets, achieves significant performance gains.
%
%% Supporting Functions
%% Mixup
% The supporting object, Mixup, is placed in your current folder when you open this example.

type Mixup

%% References
% [1] Verma, Nishchal K., et al. "Intelligent Condition Based Monitoring Using Acoustic Signals for Air Compressors." IEEE Transactions on Reliability, vol. 65, no. 1, Mar. 2016, pp. 291–309. DOI.org (Crossref), doi:10.1109/TR.2015.2459684.
% [2] Huszar, Ferenc. "Mixup: Data-Dependent Data Augmentation." InFERENCe. November 03, 2017. Accessed January 15, 2019. https://www.inference.vc/mixup-data-dependent-data-augmentation/.
% Copyright 2021 The MathWorks, Inc.