classdef ParallelTrainer < nnet.internal.cnn.Trainer
    % ParallelTrainer   Trains a network in parallel
    
    %   Copyright 2016-2022 The MathWorks, Inc.
    
    properties( Access = private )
        % InterruptStream  DataQueue created on the root worker that can be
        % used to send instructions to the pool to stop training
        InterruptStream
        
        % UseGpu  Whether to do optimized reductions on the GPU
        UseGpu
        
        % MiniBatchSize  The expected size of a full batch during training
        MiniBatchSize
        
        % MinimumMiniBatchSize  The least allowed number of observations
        % per mini-batch. Equal to MiniBatchSize for 'discardLast' end of
        % Epoch behaviour, and 1 for 'truncateLast'.
        MinimumMiniBatchSize
        
        % LeftoverObservations  Data from a previous iteration that can be
        % used to make up a full batch
        LeftoverObservations
        
        % LeftoverResponses  Responses from a previous iteration that can
        % be used to make up a full batch
        LeftoverResponses
        
        % AllNumLeftoverObs  Count of the observations in the leftovers for
        % each worker
        AllNumLeftoverObs

    end
    
    methods
        function this = ParallelTrainer( ...
                opts, precision, reporters, executionSettings, summaryFcn)
            % ParallelTrainer    Constructor for a parallel trainer
            %
            % opts - training options (nnet.cnn.TrainingOptionsSGDM)
            % precision - data precision
            % reporters - reporters for feedback during training
            % executionSettings - training environment (eg host or GPU)
            
            % Wrap reporters with ParallelReporter which knows how to call
            % back to the client from the workers
            reporters = nnet.internal.cnn.util.ParallelReporter( reporters );
            
            % Construct superclass
            this@nnet.internal.cnn.Trainer(opts, precision, reporters, executionSettings, summaryFcn);
            
            % Record execution environment
            this.UseGpu = ismember( executionSettings.executionEnvironment, {'gpu'} );
        end
        
        function net = train(this, net, data)
            % train   Train a network
            %
            % Inputs
            %    net -- network to train
            %    data -- data encapsulated in a data dispatcher
            % Outputs
            %    net -- trained network
            
            % Copy some information from the data
            this.MiniBatchSize = data.MiniBatchSize;
            this.MinimumMiniBatchSize = data.MiniBatchSize;
            if data.EndOfEpoch == "truncateLast"
                this.MinimumMiniBatchSize = 1;
            end

            % In order for SPMD to correctly understand how we wish to use
            % the Composite of data dispatchers stored within the
            % DistributedDataDispatcher, we must copy it out to a local
            % variable. We cannot access the containing dispatcher inside
            % SPMD.
            distributedData = data.DistributedData;
            
            % Determine which worker is going to have the output, and
            % create a DataQueue on that worker for the client to
            % communicate with the running trainer.
            [interruptStreamOnWorkers, this.InterruptStream] = ...
                iGetInterruptStream(distributedData);

            % Now that the communication pathway is set up the reporters
            % can be started
            this.Reporter.start();
            
            % To invoke our remote training function we need to open an
            % SPMD block. Calling SPMD is costly, so it must encompass the
            % entire algorithm.
            % In this case we cannot use the distributed dispatcher's
            % computeInParallel helper API because we need to pass a
            % Composite (the interruptStreamOnWorkers), which must be
            % declared explicitly in the same function as the SPMD block.
            trainFunc = @this.trainLocal; % Must create this outside SPMD
            spmd
                % The RemoteDispatchAdapter mixin provides the compute
                % function, which allows us to execute code treating the
                % distributed data like a normal dispatcher
                numOutputs = 2;
                [net, summary] = distributedData.compute( trainFunc, numOutputs, net, distributedData, interruptStreamOnWorkers );
                lowestRankComputeWorkerIndex = distributedData.LowestRankComputeLabIndex;
                lowestRankComputeWorkerIndex = distributedutil.AutoTransfer( lowestRankComputeWorkerIndex, lowestRankComputeWorkerIndex );
            end
            lowestRankComputeWorkerIndex = lowestRankComputeWorkerIndex.Value;
            this.Summary = summary{lowestRankComputeWorkerIndex};
            
            % net is returned as a Composite, expecting further processing
            % on the pool
        end
        
        function net = finalizeNetwork(this, net, data)
            % Always use 'truncateLast' as we want to process all the data we have.
            savedEndOfEpoch = data.EndOfEpoch;
            data.EndOfEpoch = 'truncateLast';
            
            % Accumulate
            % As for train method, we cannot use computeInParallel because
            % net is a Composite as well as data
            distributedData = data.DistributedData;
            finalizeFunc = @this.remoteFinalize;
            spmd
                net = distributedData.compute( finalizeFunc, 1, distributedData, net );
            end
            % net is returned as a Composite, expecting further processing
            % on the pool
            
            data.EndOfEpoch = savedEndOfEpoch;
        end
    end
    
    methods( Access = protected )
        
        function stopTrainingCallback(this, ~, eventData)
            % stopTraining  Overload to send cancellation request to workers
            send(this.InterruptStream, eventData);
        end
        
        function stats = computeInputStatisticsInEnvironment(this, data, stats, net )
            % Accumulate statistics in parallel
            func = @this.remoteComputeStatistics;
            [results,workerIdx] = data.computeInParallel(func,1,data,stats,net);
            % Retrieve the results from the first worker with WorkLoad>0
            results = results{workerIdx};
            % Unwrap the one output returned from func
            stats = results{1};
        end
    end
    
    methods( Access = private )
        
        function [net, summary] = trainLocal(this, net, data, interruptStream)
            % Training loop on a single worker
            trainingTimer = tic;
            prms = collectSettings(this, net);
            data.start();
            summary = this.SummaryFcn(net.OutputFormats, prms.maxEpochs);

            % Set up solve strategies
            regularizer = iCreateRegularizer('l2',net.LearnableParameters,this.Precision,this.Options);
            solver = iCreateSolver(net.LearnableParameters,this.Precision,this.Options);
            gradientThresholdOptions = struct('Method', this.Options.GradientThresholdMethod,...
                'Threshold', this.Options.GradientThreshold);
            gradThresholder = nnet.internal.cnn.GradientThresholder(gradientThresholdOptions);
            
            % Training loop - outer is epochs, inner is iterations
            iteration = 0;
            this.StopTrainingFlag = false;
            this.StopReason = nnet.internal.cnn.util.TrainingStopReason.FinalIteration;
            learnRate = initializeLearning(this);

            % Additional settings for parallel training
            prms.expectedSubBatchSizes = spmdCat(data.MiniBatchSize);
            prms.proportions = prms.expectedSubBatchSizes ./ this.MiniBatchSize;
            
            % Training loop
            for epoch = 1:prms.maxEpochs
                % For certain input data a full mini-batch can only be made
                % up by multiple reads. These arrays cache data from
                % previous reads to allow this.
                this.LeftoverObservations = [];
                this.LeftoverResponses = [];
                this.AllNumLeftoverObs = zeros(1, spmdSize);

                % Prepare data at the start of each epoch
                this.shuffle( data, prms.shuffleOption, epoch );
                data.start();
                while this.continueThisEpoch(data, interruptStream)

                    % Read a mini-batch of data. If we don't have a full
                    % mini-batch attempt to make one up from the pool of
                    % leftovers.
                    [ X, response, miniBatchSize, subBatchSize, completed ] = ...
                        this.readFullMiniBatchAttempt( data, prms );
                    if ~completed
                        % Failed to read a full mini-batch
                        continue;
                    end
                                        
                    % Allow for finished or empty datastores
                    if ~iIsAllEmpty(X)
                        % Cast data to appropriate execution environment for
                        % training
                        X = this.ExecutionStrategy.environment(X);
                        response = this.ExecutionStrategy.environment(response);
                        
                        % Compute gradients in batches. This is necessary
                        % for parallel training because uneven
                        % distributions of data can result in some workers
                        % having more observations than allowed by the
                        % training options.
                        [gradients, predictions] = this.computeGradientsInBatches( ...
                            net, X, response, subBatchSize, prms, data );
                        miniBatchLoss = net.loss( predictions, response );
                    else
                        % For finished or empty datastores, use zero
                        % gradients. This ensures these workers receive
                        % gradients even though they are not contributing.
                        subBatchSize = 0;
                        miniBatchLoss = 0;
                        gradients = iExtractZeroGradientsFromLearnableParameters( net.LearnableParameters );
                        predictions = iEmptyOrEmptyCell(prms.numOutputs);
                    end
                    
                    % Merge and normalize gradients between workers by
                    % summation
                    gradients = deep.internal.parallel.util.aggregateArrays(gradients, subBatchSize);
                    
                    % Update weights - this happens entirely on the
                    % worker, but the results are the same on every
                    % worker thus resulting in the same networks.
                    gradients = regularizer.regularizeGradients(gradients,net.LearnableParameters);
                    gradients = thresholdGradients(gradThresholder,gradients);
                    velocity = solver.calculateUpdate(gradients,learnRate);
                    net = net.updateLearnableParameters(velocity);
                    
                    % Update and report state
                    iteration = iteration + 1;
                    elapsedTime = toc(trainingTimer);
                    summary.update( predictions, response, ...
                        epoch, iteration, elapsedTime, miniBatchLoss, ...
                        learnRate, prms.lossFunctionType, data.IsDone );
                    this.Reporter.computeIteration( summary, net );
                    this.Reporter.reportIteration( summary, net );
                    this.OutputStrategy.updateOutputNetwork(summary, net);
                end
                
                learnRate = this.Schedule.update(learnRate, epoch);
                
                this.Reporter.reportEpoch( epoch, iteration, net );
                
                % If an interrupt request has been made, break out of the
                % epoch loop
                if this.StopTrainingFlag
                    break;
                end
            end  % End of epoch loop
            
            this.Reporter.computeFinish( summary, net );
            this.Reporter.finish( summary, this.StopReason );
            this.OutputStrategy.updateOutputNetwork(summary, net);
            this.OutputStrategy.addNetworkIterationToSummary(summary)
            net = this.OutputStrategy.getOutputNetwork(net);
        end
        
        function continueEpoch = continueThisEpoch(this, data, interruptStream)
            % Tests when the epoch training loop should end based on the
            % data left in the datastores across the workers, and interrupt
            % requests. These two things are set together to keep the
            % communication to a minimum. Implemented as a method so that
            % the StopTrainingFlag property can be set, to cause the epoch
            % loop to be exited as well.
            [stopRequested, stopReason, exception] = iReceivedInterrupt(interruptStream);
            stoppingConditions = ...
                struct('isDone', data.IsDone, 'stopRequested', stopRequested, 'stopReason', stopReason, 'exception', exception);
            stoppingConditions = spmdReduce(@iCatStruct, stoppingConditions);
            
            % If there were any errors, throw them collectively
            if ~isempty(stoppingConditions.exception)
                throw( stoppingConditions.exception(1) );
            end
            
            % We do not continue if any of the workers requested a stop or
            % all have finished their data AND there is no leftover data
            % for additional iterations
            numLeftOvers = sum(this.AllNumLeftoverObs);
            enoughLeftovers = numLeftOvers >= this.MinimumMiniBatchSize;
            isDone = ~enoughLeftovers && all(stoppingConditions.isDone);
            this.StopTrainingFlag = any(stoppingConditions.stopRequested);
            if ~isempty(stoppingConditions.stopReason)
                this.StopReason = stoppingConditions.stopReason(1);  
            end
            continueEpoch = ~isDone && ~this.StopTrainingFlag;
        end
        
        function [ X, response, miniBatchSize, subBatchSize, completed ] = ...
                readFullMiniBatchAttempt( this, data, prms )
            % Reads data and responses and attempt to make up a full
            % mini-batch with this and any stored leftover observations
            
            % completed indicates that enough data has been collected for
            % this iteration
            completed = true;

            % Local copy of some properties and parameters
            leftoverObservations = this.LeftoverObservations;
            leftoverResponses = this.LeftoverResponses;
            allNumLeftoverObs = this.AllNumLeftoverObs;
            observationDim = prms.InputObservationDim;
            responseObsDim = prms.OutputObservationDim;

            % Read attempt
            [X, response] = data.next();

            % Check how much data we got across the pool
            subBatchSize = iGetSizeOfFirstArrayInDim(X, observationDim);
            allSubBatchSizes = spmdCat(subBatchSize);
            miniBatchSize = sum(allSubBatchSizes);
            subBatchSize = allSubBatchSizes(spmdIndex);
            
            % If we have enough data, return now
            if miniBatchSize == this.MiniBatchSize
                return
            end
            
            % In the special circumstance where we are on the last
            % mini-batch and have truncate behaviour, return all the
            % leftovers to be processed as the final batch
            numLeftOverObs = sum(allNumLeftoverObs);
            if miniBatchSize == 0 && numLeftOverObs >= this.MinimumMiniBatchSize
                X = leftoverObservations;
                response = leftoverResponses;
                this.LeftoverObservations = [];
                this.LeftoverResponses = [];
                this.AllNumLeftoverObs = zeros(1, spmdSize);
                return
            end

            % If we don't have enough new data + leftovers, add to the
            % data and keep reading
            numMissingObs = this.MiniBatchSize - miniBatchSize;
            if numLeftOverObs < numMissingObs

                [this.LeftoverObservations,this.LeftoverResponses] = data.CatAndSliceStrategy.catBatch(...
                    observationDim,responseObsDim,X,leftoverObservations,response,leftoverResponses);

                this.AllNumLeftoverObs = allSubBatchSizes + allNumLeftoverObs;
                completed = false;
                return
            end
            
            % We have enough leftovers to create a full
            % mini-batch
            
            % First, attempt to make every worker up to a full
            % sub-batch
            myNumLeftoverObservations = allNumLeftoverObs(spmdIndex);
            allNumMissingObs = prms.expectedSubBatchSizes - allSubBatchSizes;
            allAddedObservations1 = min(allNumMissingObs, allNumLeftoverObs);
            allNumLeftoverObs = allNumLeftoverObs - allAddedObservations1;
            numAddedObs = sum(allAddedObservations1);
            numMissingObs = numMissingObs - numAddedObs;
            
            % Now attempt to assign remaining leftovers,
            % respecting the original choice of distribution
            adjustedProportions = prms.proportions;
            adjustedProportions(allNumLeftoverObs==0) = 0;
            adjustedProportions = adjustedProportions ./ sum(adjustedProportions);
            allAddedObservations2 = min(allNumLeftoverObs, floor(adjustedProportions.*numMissingObs));
            allNumLeftoverObs = allNumLeftoverObs - allAddedObservations2;
            numMissingObs = numMissingObs - sum(allAddedObservations2);
            
            % Finally, just assign any remaining leftovers
            % evenly
            numComputeWorkers = numel(allNumLeftoverObs);
            allAddedObservations3 = zeros(max(allNumLeftoverObs), numComputeWorkers);
            for w = 1:numComputeWorkers
                allAddedObservations3(1:allNumLeftoverObs(w), w) = 1;
            end
            % Find point where we reached a full mini-batch
            count = cumsum(allAddedObservations3(:));
            allAddedObservations3(count>numMissingObs) = 0;
            allAddedObservations3 = sum(allAddedObservations3, 1);
            allNumLeftoverObs = allNumLeftoverObs - allAddedObservations3;
            
            % Add the appropriate number of leftover
            % observations to the data
            allAddedObservations = allAddedObservations1 + allAddedObservations2 + allAddedObservations3;
            myAddedObservations = allAddedObservations(spmdIndex);
            trimdices = 1:myAddedObservations;

            [extrasX, extrasResponse] = data.CatAndSliceStrategy.sliceBatch(...
                observationDim,responseObsDim,trimdices,leftoverObservations,leftoverResponses);
            [X,response] = data.CatAndSliceStrategy.catBatch(...
                    observationDim,responseObsDim,X,extrasX,response,extrasResponse);

            % Trim the remaining leftovers
            trimdices = (myAddedObservations+1):myNumLeftoverObservations;

            [this.LeftoverObservations,this.LeftoverResponses] = data.CatAndSliceStrategy.sliceBatch(...
                observationDim,responseObsDim,trimdices,leftoverObservations,leftoverResponses);

            miniBatchSize = this.MiniBatchSize;
            allSubBatchSizes = allSubBatchSizes + myAddedObservations;
            subBatchSize = allSubBatchSizes(spmdIndex);
            this.AllNumLeftoverObs = allNumLeftoverObs;
        end
        
        function [gradients, predictions] = computeGradientsInBatches(this, ...
                net, X, response, subBatchSize, prms, data)
            % Compute gradients, but loop to restrict the batch size to
            % that expected by the training options
            unitSize = prms.expectedSubBatchSizes(spmdIndex);
            if subBatchSize <= unitSize
                % Normal case where we got a full sub-batch, avoid some
                % unnecessary extra work
                [gradients, predictions] = this.computeGradients(net, ...
                    X, response);
            else
                % Overfull sub-batch: sum gradients and concatenate
                % predictions
                gradients = iExtractZeroGradientsFromLearnableParameters( net.LearnableParameters );
                predictions = iEmptyOrEmptyCell(prms.numOutputs);
                for obsIndex = 1:unitSize:subBatchSize
                    % Get slice of the data
                    endOfRange = min(subBatchSize,(obsIndex+unitSize-1));
                    indices = obsIndex:endOfRange;
                    [Xpart,responsePart] = data.CatAndSliceStrategy.sliceBatch(...
                        prms.InputObservationDim,prms.OutputObservationDim,indices,X,response);
                    % Compute gradients and predictions
                    [gradientsPart, predictionsPart] = this.computeGradients(net, Xpart, responsePart);
                    % Aggregate results appropriately
                    gradients = cellfun(@(g1,g2)g1+g2, gradients, gradientsPart, 'UniformOutput', false);
                    predictions = nnet.internal.cnn.util.NDArrayCatAndSliceStrategy.catCellBatch(prms.OutputObservationDim, predictions, predictionsPart);
                end
            end
        end
        
        function stats = remoteComputeStatistics(this, distData, stats, net )
            % Initialization steps running on the pool workers

            % Perform actual initialization with data
            stats = this.doComputeStatistics(distData,stats,net);
            
            % Merge the distributed statistics from different workers and
            % accumulate the results onto the first compute worker
            for i = 1:numel(stats)
                stats{i} = spmdReduce( @merge, stats{i}, 1 );
            end
            
            % Statistics returned should be ready for use on the client,
            % which may not have a GPU
            if spmdIndex == 1
                stats = cellfun( @gather, stats, 'UniformOutput', false);
            end
        end
        
        function net = remoteFinalize(this, distData, net)
            % Finalize algorithm running on the pool workers
            
            % Make sure data is in right place
            if this.UseGpu
                net = net.setupNetworkForGPUTraining(true);
            else
                net = net.setupNetworkForHostTraining();
            end
            net = this.doFinalize(net, distData);
            
            % Merge all finalizable layers, putting the result on
            % the first worker in the split communicator.
            for ii = 1:numel(net.Layers)
                thisLayer = net.Layers{ii};
                if isa(thisLayer,'nnet.internal.cnn.layer.Finalizable') && ...
                        thisLayer.NeedsFinalize
                    net.Layers{ii} = spmdReduce( @mergeFinalized, thisLayer, 1 );
                end
            end
            
            % Network returned should be ready for use on the client,
            % which may not have a GPU
            if spmdIndex == 1
                net = net.setupNetworkForHostTraining();
            end
        end
        
    end
    
end

function [interruptStreamWorkers, interruptStreamClient] = ...
    iGetInterruptStream( distributedData )
% Determines which worker will hold the output network during training, and
% creates a DataQueue on that worker for the client to communicate with to
% control training.
spmd
    % The worker with the result will be the compute worker with the lowest
    % rank
    isWorkerActive = distributedData.IsComputeWorker;
    workerIndexWithOutput = spmdIndex;
    if ~isWorkerActive
        workerIndexWithOutput = inf;
    end
    workerIndexWithOutput = spmdReduce(@min, workerIndexWithOutput);
    
    % Create a DataQueue on this worker and return a Composite
    interruptStreamWorkers = parallel.internal.pool.DataQueue.empty;
    if spmdIndex == workerIndexWithOutput
        interruptStreamWorkers = parallel.internal.pool.DataQueue;
    end
    
    % Return results to client
    interruptStreamClient = distributedutil.AutoTransfer( interruptStreamWorkers, workerIndexWithOutput );
end

% Retrieve underlying data
interruptStreamClient = interruptStreamClient.Value;
end

function [stopRequested, stopReason, exception] = iReceivedInterrupt( interruptStream )
% Checks the data queue on the root worker to see if an interrupt
% is being requested
stopRequested = false;
stopReason = [];
exception = [];
if ~isempty(interruptStream)
    [data, ok] = poll(interruptStream);
    if ok
        stopRequested = true;
        if isa( data, 'nnet.internal.cnn.util.TrainingInterruptEventData' )
            stopReason = data.TrainingStopReason;
        elseif isa( data, 'nnet.internal.cnn.util.ParallelTrainingInterruptEventData' )
            exception = data.Exception; 
        end
    end
end
end

function regularizer = iCreateRegularizer(name,learnableParameters,precision,regularizationOptions)
regularizer = nnet.internal.cnn.regularizer.RegularizerFactory.create(name,learnableParameters,precision,regularizationOptions);
end

function solver = iCreateSolver(learnableParameters,precision,trainingOptions)
solver = nnet.internal.cnn.solver.SolverFactory.create(learnableParameters,precision,trainingOptions);
end

function zeroGradients = iExtractZeroGradientsFromLearnableParameters(learnableParametersArray)
zeroGradients = cell(1,numel(learnableParametersArray));
for i = 1:numel(learnableParametersArray)
    thisParam = learnableParametersArray(i).Value;
    zeroGradients{i} = zeros( size(thisParam), 'like', thisParam );
end
end

function s = iCatStruct(s1, s2)
% Concatenate the fields of two structures
s = s1;
f = fieldnames(s);
for i = 1:numel(f)
    fname = f{i};
    if ischar(s1.(fname))
        s.(fname) = {s1.(fname); s2.(fname)};
    else
        s.(fname) = [s1.(fname); s2.(fname)];
    end
end
end

function sz = iGetSizeInDim(X, dim)
% Returns size(X,dim) except when X is empty in which case always return 0
sz = 0;
if ~isempty(X)
    sz = size(X, dim);
end
end

function sz = iGetSizeOfFirstArrayInDim(X, dim)
% Returns the size of X or X{1} in dim
if iscell(X)
    sz = iGetSizeInDim(X{1}, dim{1});
else
    sz = iGetSizeInDim(X, iUnWrapScalarCell(dim));
end
end

function tf = iIsAllEmpty(X)
tf = isempty(X) || iscell(X) && any(cellfun(@isempty,X));
end

function X = iUnWrapScalarCell(X)
if iscell(X) && numel(X) == 1
    X = X{1};
end
end

function X = iEmptyOrEmptyCell(nelements)
if nelements == 1
    X = [];
else
    X = cell(1,nelements);
end
end
