classdef Trainer < handle
    %Trainer   Trains a network

    %   Copyright 2015-2022 The MathWorks, Inc.

    properties(Access = protected)
        Options
        Schedule
        Precision
        Reporter
        SummaryFcn
        Summary
        ExecutionStrategy
        StopTrainingFlag
        StopReason
        InterruptException
        OutputStrategy
    end

    methods
        function this = Trainer( ...
                opts, precision, reporter, executionSettings, summaryFcn)
            % Trainer    Constructor for a network trainer
            %
            % opts - training options (nnet.cnn.TrainingOptionsSGDM)
            % precision - data precision
            this.Options = opts;
            this.Schedule = nnet.internal.cnn.LearnRateScheduleFactory.create(opts.LearnRateScheduleSettings);
            this.Precision = precision;
            this.Reporter = reporter;
            this.SummaryFcn = summaryFcn;
            this.OutputStrategy = nnet.internal.cnn.trainNetwork.outputnetwork.OutputNetworkStrategyFactory.createStrategy(opts.OutputNetwork, executionSettings);
            % Declare execution strategy
            if ismember( executionSettings.executionEnvironment, {'gpu'} )
                this.ExecutionStrategy = nnet.internal.cnn.TrainerGPUStrategy;
            else
                this.ExecutionStrategy = nnet.internal.cnn.TrainerHostStrategy;
            end

            % Print execution environment if in verbose mode
            iPrintExecutionEnvironment(opts, executionSettings);

            % Register a listener to detect requests to terminate training
            addlistener( reporter, ...
                'TrainingInterruptEvent', @this.stopTrainingCallback);
        end

        function net = initializeNetwork(this, net, data)
            % Perform any initialization steps required by the input and
            % output layers

            % Store any classification labels or response names in the
            % appropriate output layers.
            net = net.storeResponseData(data.ResponseMetaData);

            if this.Options.ResetInputNormalization
                net = net.resetNetworkInitialization();
            end

            isMovingAverage = isequal(this.Options.BatchNormalizationStatistics,'moving');
            net = net.setBatchNormalizationStatisticsMode(isMovingAverage);

            % Setup reporters
            willInputNormalizationBeComputed = net.hasEmptyInputStatistics();
            reporterSetupInfo = nnet.internal.cnn.util.ReporterSetupInfo(willInputNormalizationBeComputed);
            this.Reporter.setup(reporterSetupInfo);

            if willInputNormalizationBeComputed
                % Always use 'truncateLast' as we want to process all the data we have.
                savedEndOfEpoch = data.EndOfEpoch;
                data.EndOfEpoch = 'truncateLast';

                % Compute statistics on host or in distributed environment
                stats = iCreateStatisticsAccumulator(net.InputSizes,this.Precision);
                stats = this.computeInputStatisticsInEnvironment(data,stats,net);

                % Initialize input layers with statistics
                net = net.initializeNetwork(stats);

                data.EndOfEpoch = savedEndOfEpoch;
            end
        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
            reporter = this.Reporter;
            schedule = this.Schedule;
            prms = collectSettings(this, net);
            outputStrategy = this.OutputStrategy;
            data.start();
            this.Summary = this.SummaryFcn(net.OutputFormats, prms.maxEpochs);

            regularizer = iCreateRegularizer('l2',net.LearnableParameters,this.Precision,this.Options);

            solver = iCreateSolver(net.LearnableParameters,this.Precision,this.Options);

            trainingTimer = tic;

            % Set StopTrainingFlag and StopReason before starting reporters
            % since they might stop training
            this.StopTrainingFlag = false;
            this.StopReason = nnet.internal.cnn.util.TrainingStopReason.FinalIteration;
            reporter.start();
            iteration = 0;
            gradientThresholdOptions = struct('Method', this.Options.GradientThresholdMethod,...
                'Threshold', this.Options.GradientThreshold);
            gradThresholder = nnet.internal.cnn.GradientThresholder(gradientThresholdOptions);
            learnRate = initializeLearning(this);

            for epoch = 1:prms.maxEpochs
                this.shuffle( data, prms.shuffleOption, epoch );
                data.start();
                while ~data.IsDone && ~this.StopTrainingFlag
                    [X, T, info] = data.next();
                    % Cast data to appropriate execution environment for
                    % training and apply transforms
                    X = this.ExecutionStrategy.environment(X);
                    T = this.ExecutionStrategy.environment(T);

                    [gradients, Y, states] = this.computeGradients(net, X, T);

                    % Reuse the layers outputs to compute loss
                    miniBatchLoss = net.loss( Y, T );

                    gradients = regularizer.regularizeGradients(gradients,net.LearnableParameters);

                    gradients = thresholdGradients(gradThresholder,gradients);

                    velocity = solver.calculateUpdate(gradients,learnRate);

                    net = net.updateLearnableParameters(velocity);

                    net = this.setNetworkState(net, states, info.IsCompleteObservation);

                    elapsedTime = toc(trainingTimer);

                    iteration = iteration + 1;
                    this.Summary.update(Y, T, epoch, iteration, elapsedTime, miniBatchLoss, learnRate, data.IsDone);
                    % It is important that computeIteration is called
                    % before reportIteration, so that the summary is
                    % correctly updated before being reported
                    reporter.computeIteration( this.Summary, net );
                    reporter.reportIteration( this.Summary, net );

                    outputStrategy.updateOutputNetwork(this.Summary, net);
                end
                learnRate = schedule.update(learnRate, epoch);

                reporter.reportEpoch( epoch, iteration, net );

                % If an interrupt request has been made, break out of the
                % epoch loop
                if this.StopTrainingFlag
                    % If the summary has not been updated, set its
                    % Iteration count to zero before breaking out of the
                    % epoch loop. This ensures subsequent Reporter methods
                    % can be called.
                    if isempty(this.Summary.Iteration)
                        this.Summary.Iteration = 0;
                    end

                    break;
                end
            end
            reporter.computeFinish( this.Summary, net );
            reporter.finish( this.Summary, this.StopReason );
            % Validation is always computed at the end of training if it
            % hasn't been computed in the last iteration. Update output
            % network before getting output network
            outputStrategy.updateOutputNetwork(this.Summary, net);
            net = outputStrategy.getOutputNetwork( net );
            % Add the iteration corresponding to net to summary so
            % that any reporter that needs the output network iteration at
            % reportFinalIteration will have access to it
            outputStrategy.addNetworkIterationToSummary( this.Summary )
        end

        function net = finalizeNetwork(this, net, data)
            % Perform any finalization steps required by the layers

            % Always use 'truncateLast' as we want to process all the data we have.
            savedEndOfEpoch = data.EndOfEpoch;
            data.EndOfEpoch = 'truncateLast';

            % Call shared implementation
            net = this.doFinalize(net, data);

            data.EndOfEpoch = savedEndOfEpoch;
        end
    end

    methods(Access = protected)
        function stopTrainingCallback(this, ~, evtData)
            % stopTraining  Callback triggered by interrupt events that
            % want to request training to stop
            this.StopTrainingFlag = true;
            this.StopReason = evtData.TrainingStopReason;
        end

        function settings = collectSettings(this, net)
            % collectSettings  Collect together fixed settings from the
            % Trainer and the data and put in the correct form.
            settings.maxEpochs = this.Options.MaxEpochs;
            settings.lossFunctionType = iGetLossFunctionType(net);
            settings.shuffleOption = this.Options.Shuffle;
            settings.numOutputs = numel(net.OutputSizes);
            settings.InputObservationDim = cellfun(@(sz)numel(sz)+1, net.InputSizes, 'UniformOutput', false);
            settings.OutputObservationDim = cellfun(@(sz)numel(sz)+1, net.OutputSizes, 'UniformOutput', false);
        end

        function learnRate = initializeLearning(this)
            % initializeLearning  Set initial learning rate.
            learnRate = this.Precision.cast( this.Options.InitialLearnRate );
        end

        function [gradients, predictions, states] = computeGradients(~, net, X, Y)
            % computeGradients   Compute the gradients of the network. This
            % function returns also the network output so that we will not
            % need to perform the forward propagation step again.
            propagateState = true;
            [gradients, predictions, states] = net.computeGradientsForTraining(X, Y, propagateState);
        end

        function net = setNetworkState(~, net, states, needToResetState)
            % setNetworkState   Set the network state. First update the
            % state based on the states computed during the forward pass,
            % then we reset if necessary.
            net = net.updateNetworkState(states);
            if needToResetState
                net = net.resetNetworkState();
            end
        end

        function stats = computeInputStatisticsInEnvironment(this, data, stats, net )
            % Initialization steps running on the host
            stats = doComputeStatistics(this, data, stats, net);
        end

        function stats = doComputeStatistics(this, data, stats, net)
            % Compute statistics from the training data set
            % Do one epoch

            % If we have a sequence dispatcher which applies padding we
            % replace padding values with NaN. Then padding can be ignored
            % during computation of statistics
            replacePaddingWithNaN = iNeedToReplacePaddingWithNaNs(data);

            data.start();
            while ~data.IsDone
                if replacePaddingWithNaN
                    [X, ~, info] = data.next();
                    % Replace padded values with NaN, since NaNs will be
                    % ignored when computing statistics
                    mask = info.Mask;
                    if any(~mask, "all")
                        % Permute the mask from (S)x...x(C)x(T)x(B) to
                        % (S)x...x(C)x(B)x(T)
                        numSCDims = numel(data.DataSize);
                        mask = permute(mask, [1:numSCDims numSCDims+2 numSCDims+1]);
                        mask = double(mask);
                        mask(~mask) = NaN;
                        X = mask.*X;
                    end
                else
                    X = data.next();
                end
                if ~isempty(X) % In the parallel case X can be empty
                    % Cast data to appropriate execution environment
                    X = this.ExecutionStrategy.environment(X);
                    X = iWrapInCell(X);

                    % Accumulate statistics for each input
                    for i = 1:numel(stats)
                        % The ImageInputLayer might have data augmentation
                        % transforms that have to be applied before
                        % computing statistics (e.g. cropping).
                        if isa(net.InputLayers{i},'nnet.internal.cnn.layer.ImageInput')
                            X{i} = apply(net.InputLayers{i}.TrainTransforms,X{i});
                        end

                        stats{i} = accumulate(stats{i}, X{i});
                    end
                end
            end
        end

        function net = doFinalize(this, net, data)
            % Perform any finalization steps required by the layers
            isFinalizable = @(l)isa(l,'nnet.internal.cnn.layer.Finalizable') && ...
                l.NeedsFinalize;
            needsFinalize = cellfun(isFinalizable, net.Layers);
            if any(needsFinalize)
                % Do one final epoch
                data.start();
                while ~data.IsDone
                    X = data.next();
                    if ~isempty(X) % In the parallel case X can be empty
                        % Cast data to appropriate execution environment for
                        % training and apply transforms
                        X = this.ExecutionStrategy.environment(X);

                        % Ask the network to finalize
                        net = finalizeNetwork(net, X);
                    end
                end
            end
            this.Reporter.computeFinalIteration(this.Summary, net);
            this.Reporter.reportFinalIteration(this.Summary);
        end
    end

    methods(Access = protected)
        function shuffle(~, data, shuffleOption, epoch)
            % shuffle   Shuffle the data as per training options
            if ~isequal(shuffleOption, 'never') && ...
                    ( epoch == 1 || isequal(shuffleOption, 'every-epoch') )
                data.shuffle();
            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 stats = iCreateStatisticsAccumulator(inputSizes,outputType)
stats = nnet.internal.cnn.statistics.AccumulatorFactory.create(inputSizes,outputType);
end

function tf = iNeedToReplacePaddingWithNaNs(data)
tf = isa(data,'nnet.internal.cnn.dispatcher.BuiltInSequenceDispatcher') ...
    && ~isequal(data.SequenceLength,'shortest') ...
    && ~isnan(data.PaddingValue);
end

function t = iGetLossFunctionType(net)
if isempty(net.Layers)
    t = 'nnet.internal.cnn.layer.NullLayer';
else
    t = class(net.Layers{end});
end
end

function iPrintMessage(messageID, varargin)
string = getString(message(messageID, varargin{:}));
fprintf( '%s\n', string );
end

function iPrintExecutionEnvironment(opts, executionSettings)
% Print execution environment if in 'auto' mode
if opts.Verbose
    if ismember(opts.ExecutionEnvironment, {'auto'})
        if ismember(executionSettings.executionEnvironment, {'cpu'})
            iPrintMessage( ...
                'nnet_cnn:internal:cnn:Trainer:TrainingInSerialOnCPU');
        elseif ismember(executionSettings.executionEnvironment, {'gpu'})
            if executionSettings.useStateless
                iPrintMessage( ...
                    'nnet_cnn:internal:cnn:Trainer:TrainingStatelessInSerialOnGPU');
            else
                iPrintMessage( ...
                    'nnet_cnn:internal:cnn:Trainer:TrainingInSerialOnGPU');
            end
        end
    elseif ismember(opts.ExecutionEnvironment, {'parallel'})
        if ismember(executionSettings.executionEnvironment, {'cpu'})
            iPrintMessage( ...
                'nnet_cnn:internal:cnn:Trainer:TrainingInParallelOnCPUs');
        elseif ismember(executionSettings.executionEnvironment, {'gpu'})
            iPrintMessage( ...
                'nnet_cnn:internal:cnn:Trainer:TrainingInParallelOnGPUs');
        end
    end
end
end

function X = iWrapInCell(X)
if ~iscell(X)
    X = {X};
end
end
