vllab_cnn_train_dagNN_multi_label.m

function [net,stats] = vllab_cnn_train_dagNN_multi_label(net, imdb,  varargin)
% -------------------------------------------------------------------------
%   Description:
%       Modified script for training SDCNN with dagnn
%
%   Input:
%       - net       : dagnn model
%       - imdb      : training dataset
%
%   Citation:
%       Deep Laplacian Pyramid Networks for Fast and Accurate Super-Resolution
%       Wei-Sheng Lai, Jia-Bin Huang, Narendra Ahuja, and Ming-Hsuan Yang
%       IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2017
%
%   Contact:
%       Wei-Sheng Lai
%       wlai24@ucmerced.edu
%       University of California, Merced
% -------------------------------------------------------------------------

%CNN_TRAIN_DAG Demonstrates training a CNN using the DagNN wrapper
%    CNN_TRAIN_DAG() is similar to CNN_TRAIN(), but works with
%    the DagNN wrapper instead of the SimpleNN wrapper.

% Copyright (C) 2014-16 Andrea Vedaldi.
% All rights reserved.
%
% This file is part of the VLFeat library and is made available under
% the terms of the BSD license (see the COPYING file).
addpath(fullfile(vl_rootnn, 'examples'));

opts.expDir = fullfile('data','exp') ;
opts.continue = true ;
opts.batchSize = 256 ;
opts.numSubBatches = 1 ;
opts.train = [] ;
opts.val = [] ;
opts.gpus = [] ;
opts.prefetch = false ;
opts.epochSize = inf;
opts.numEpochs = 300 ;
opts.learningRate = 0.001 ;
opts.weightDecay = 0.0005 ;

opts.solver = [] ;  % Empty array means use the default SGD solver
[opts, varargin] = vl_argparse(opts, varargin) ;
if ~isempty(opts.solver)
    assert(isa(opts.solver, 'function_handle') && nargout(opts.solver) == 2,...
        'Invalid solver; expected a function handle with two outputs.') ;
    % Call without input arguments, to get default options
    opts.solverOpts = opts.solver() ;
end

opts.momentum = 0.9 ;
opts.saveSolverState = true ;
opts.nesterovUpdate = false ;
opts.randomSeed = 0 ;
opts.profile = false ;
opts.parameterServer.method = 'mmap' ;
opts.parameterServer.prefix = 'mcn' ;

opts.derOutputs = {'objective', 1} ;
opts.extractStatsFn = @extractStats ;
opts.plotStatistics = true;
opts.postEpochFn = [] ;  % postEpochFn(net,params,state) called after each epoch; can return a new learning rate, 0 to stop, [] for no change

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% additional opts
opts.model_name             = '';
opts.num_train_batch        = 1000;
opts.num_valid_batch        = 100;
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%


opts = vl_argparse(opts, varargin) ;

if ~exist(opts.expDir, 'dir'), mkdir(opts.expDir) ; end
if isempty(opts.train), opts.train = find(imdb.set==1) ; end
if isempty(opts.val), opts.val = find(imdb.set==2) ; end
if isscalar(opts.train) && isnumeric(opts.train) && isnan(opts.train)
    opts.train = [] ;
end
if isscalar(opts.val) && isnumeric(opts.val) && isnan(opts.val)
    opts.val = [] ;
end

% -------------------------------------------------------------------------
%                                                            Initialization
% -------------------------------------------------------------------------

evaluateMode = isempty(opts.train) ;
if ~evaluateMode
    if isempty(opts.derOutputs)
        error('DEROUTPUTS must be specified when training.\n') ;
    end
end

% -------------------------------------------------------------------------
%                                                        Train and validate
% -------------------------------------------------------------------------

modelPath = @(ep) fullfile(opts.expDir, sprintf('net-epoch-%d.mat', ep));
modelFigPath = fullfile(opts.expDir, 'net-train.pdf') ;

start = opts.continue * findLastCheckpoint(opts.expDir) ;
if start >= 1
    fprintf('%s: resuming by loading epoch %d\n', mfilename, start) ;
    [net, state, stats] = loadState(modelPath(start)) ;
else
    state = [] ;
end

for epoch=start+1:opts.numEpochs
    
    % Set the random seed based on the epoch and opts.randomSeed.
    % This is important for reproducibility, including when training
    % is restarted from a checkpoint.
    
    rng(epoch + opts.randomSeed) ;
    prepareGPUs(opts, epoch == start+1) ;
    
    % Train for one epoch.
    params = opts ;
    params.epoch = epoch ;
    params.learningRate = opts.learningRate(min(epoch, numel(opts.learningRate))) ;
    params.train = opts.train(randperm(numel(opts.train))) ; % shuffle
    params.train = params.train(1:min(opts.epochSize, numel(opts.train)));
    params.val = opts.val(randperm(numel(opts.val))) ;
    params.imdb = imdb ;
    params.getBatch = getBatch ;
    
    if numel(opts.gpus) <= 1
        [net, state] = processEpoch(net, state, params, 'train') ;
        [net, state] = processEpoch(net, state, params, 'val') ;
        if ~evaluateMode
            saveState(modelPath(epoch), net, state) ;
        end
        lastStats = state.stats ;
    else
        spmd
            [net, state] = processEpoch(net, state, params, 'train') ;
            [net, state] = processEpoch(net, state, params, 'val') ;
            if labindex == 1 && ~evaluateMode
                saveState(modelPath(epoch), net, state) ;
            end
            lastStats = state.stats ;
        end
        lastStats = accumulateStats(lastStats) ;
    end
    
    stats.train(epoch) = lastStats.train ;
    stats.val(epoch) = lastStats.val ;
    clear lastStats ;
    saveStats(modelPath(epoch), stats) ;
    
    if opts.plotStatistics
        switchFigure(1) ; clf ;
        plots = setdiff(...
            cat(2,...
            fieldnames(stats.train)', ...
            fieldnames(stats.val)'), {'num', 'time'}) ;
        for p = plots
            p = char(p) ;
            values = zeros(0, epoch) ;
            leg = {} ;
            for f = {'train', 'val'}
                f = char(f) ;
                if isfield(stats.(f), p)
                    tmp = [stats.(f).(p)] ;
                    values(end+1,:) = tmp(1,:)' ;
                    leg{end+1} = f ;
                end
            end
            subplot(1,numel(plots),find(strcmp(p,plots))) ;
            plot(1:epoch, values','o-') ;
            xlabel('epoch') ;
            %       title(p) ;
            title(strrep(p, '_', '-')) ;
            legend(leg{:}) ;
            grid on ;
        end
        drawnow ;
        print(1, modelFigPath, '-dpdf') ;
    end
    
    if ~isempty(opts.postEpochFn)
        if nargout(opts.postEpochFn) == 0
            opts.postEpochFn(net, params, state) ;
        else
            lr = opts.postEpochFn(net, params, state) ;
            if ~isempty(lr), opts.learningRate = lr; end
            if opts.learningRate == 0, break; end
        end
    end
end

% With multiple GPUs, return one copy
if isa(net, 'Composite'), net = net{1} ; end

% -------------------------------------------------------------------------
function [net, state] = processEpoch(net, state, params, mode)
% -------------------------------------------------------------------------
% Note that net is not strictly needed as an output argument as net
% is a handle class. However, this fixes some aliasing issue in the
% spmd caller.

% initialize with momentum 0
if isempty(state) || isempty(state.solverState)
    state.solverState = cell(1, numel(net.params)) ;
    state.solverState(:) = {0} ;
end

% move CNN  to GPU as needed
numGpus = numel(params.gpus) ;
if numGpus >= 1
    net.move('gpu') ;
    for i = 1:numel(state.solverState)
        s = state.solverState{i} ;
        if isnumeric(s)
            state.solverState{i} = gpuArray(s) ;
        elseif isstruct(s)
            state.solverState{i} = structfun(@gpuArray, s, 'UniformOutput', false) ;
        end
    end
end
if numGpus > 1
    parserv = ParameterServer(params.parameterServer) ;
    net.setParameterServer(parserv) ;
else
    parserv = [] ;
end

% profile
if params.profile
    if numGpus <= 1
        profile clear ;
        profile on ;
    else
        mpiprofile reset ;
        mpiprofile on ;
    end
end

num = 0 ;
epoch = params.epoch ;
subset = params.(mode) ;

% %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% image_set = params.(mode) ;
%
% if( strcmp(mode,'train') )
%     num_samples = params.num_train_batch * params.batchSize;
% elseif( strcmp(mode,'val') )
%     num_samples = params.num_valid_batch * params.batchSize;
% end
%
% subset = [];
% while( length(subset) < num_samples )
%     subset = cat(1, subset, image_set(:));
% end
% subset = subset(1:num_samples);
%
% %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%




adjustTime = 0 ;

stats.num = 0 ; % return something even if subset = []
stats.time = 0 ;

start = tic ;
for t=1:params.batchSize:numel(subset)
    %   fprintf('%s: epoch %02d: %3d/%3d:', mode, epoch, ...
    %           fix((t-1)/params.batchSize)+1, ceil(numel(subset)/params.batchSize)) ;
    batchSize = min(params.batchSize, numel(subset) - t + 1) ;
    
    for s=1:params.numSubBatches
        % get this image batch and prefetch the next
        batchStart = t + (labindex-1) + (s-1) * numlabs ;
        batchEnd = min(t+params.batchSize-1, numel(subset)) ;
        batch = subset(batchStart : params.numSubBatches * numlabs : batchEnd) ;
        num = num + numel(batch) ;
        if numel(batch) == 0, continue ; end
        
        inputs = params.getBatch(params,params.imdb, batch) ;
        
        if params.prefetch
            if s == params.numSubBatches
                batchStart = t + (labindex-1) + params.batchSize ;
                batchEnd = min(t+2*params.batchSize-1, numel(subset)) ;
            else
                batchStart = batchStart + numlabs ;
            end
            nextBatch = subset(batchStart : params.numSubBatches * numlabs : batchEnd) ;
            params.getBatch(params.imdb, nextBatch) ;
        end
        
        if strcmp(mode, 'train')
            net.mode = 'normal' ;
            net.accumulateParamDers = (s ~= 1) ;
            net.eval(inputs, params.derOutputs, 'holdOn', s < params.numSubBatches) ;
        else
            net.mode = 'test' ;
            net.eval(inputs) ;
        end
    end
    
    % Accumulate gradient.
    if strcmp(mode, 'train')
        if ~isempty(parserv), parserv.sync() ; end
        state = accumulateGradients(net, state, params, batchSize, parserv) ;
    end
    
    % Get statistics.
    time = toc(start) + adjustTime ;
    batchTime = time - stats.time ;
    stats.num = num ;
    stats.time = time ;
    stats = params.extractStatsFn(stats,net) ;
    currentSpeed = batchSize / batchTime ;
    averageSpeed = (t + batchSize - 1) / time ;
    if t == 3*params.batchSize + 1
        % compensate for the first three iterations, which are outliers
        adjustTime = 4*batchTime - time ;
        stats.time = time + adjustTime ;
    end
    
    %   fprintf(' %.1f (%.1f) Hz', averageSpeed, currentSpeed) ;
    
    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    if( numel(params.gpus) > 0 )
        fprintf('%s [gpu %d]: ', mode, params.gpus(1)) ;
    else
        fprintf('%s [cpu]: ', mode) ;
    end
    fprintf('epoch %02d: %3d/%3d: %.1f (%.1f) Hz, lr = %s\n', ...
        epoch, ...
        fix((t-1)/params.batchSize)+1, ...
        ceil(numel(subset)/params.batchSize),...
        averageSpeed, currentSpeed, ...
        num2str(params.learningRate));
    
    fprintf('\tModel = %s\n', params.model_name);
    
    for f = setdiff(fieldnames(stats)', {'num', 'time'})
        fc = char(f) ;
        fprintf('\t%s: %.4f\n', fc, stats.(fc)) ;
    end
    
    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    
end

% Save back to state.
state.stats.(mode) = stats ;
if params.profile
    if numGpus <= 1
        state.prof.(mode) = profile('info') ;
        profile off ;
    else
        state.prof.(mode) = mpiprofile('info');
        mpiprofile off ;
    end
end
if ~params.saveSolverState
    state.solverState = [] ;
else
    for i = 1:numel(state.solverState)
        s = state.solverState{i} ;
        if isnumeric(s)
            state.solverState{i} = gather(s) ;
        elseif isstruct(s)
            state.solverState{i} = structfun(@gather, s, 'UniformOutput', false) ;
        end
    end
end

net.reset() ;
net.move('cpu') ;

% -------------------------------------------------------------------------
function state = accumulateGradients(net, state, params, batchSize, parserv)
% -------------------------------------------------------------------------
numGpus = numel(params.gpus) ;
otherGpus = setdiff(1:numGpus, labindex) ;

for p=1:numel(net.params)
    
    if ~isempty(parserv)
        parDer = parserv.pullWithIndex(p) ;
    else
        parDer = net.params(p).der ;
    end
    
    switch net.params(p).trainMethod
        
        case 'average' % mainly for batch normalization
            thisLR = net.params(p).learningRate ;
            net.params(p).value = vl_taccum(...
                1 - thisLR, net.params(p).value, ...
                (thisLR/batchSize/net.params(p).fanout),  parDer) ;
            
        case 'gradient'
            thisDecay = params.weightDecay * net.params(p).weightDecay ;
            thisLR = params.learningRate * net.params(p).learningRate ;
            
            if thisLR>0 || thisDecay>0
                % Normalize gradient and incorporate weight decay.
                parDer = vl_taccum(1/batchSize, parDer, ...
                    thisDecay, net.params(p).value) ;
                
                if isempty(params.solver)
                    % Default solver is the optimised SGD.
                    % Update momentum.
                    state.solverState{p} = vl_taccum(...
                        params.momentum, state.solverState{p}, ...
                        -1, parDer) ;
                    
                    % Nesterov update (aka one step ahead).
                    if params.nesterovUpdate
                        delta = params.momentum * state.solverState{p} - parDer ;
                    else
                        delta = state.solverState{p} ;
                    end
                    
                    % Update parameters.
                    net.params(p).value = vl_taccum(...
                        1,  net.params(p).value, thisLR, delta) ;
                    
                else
                    % call solver function to update weights
                    [net.params(p).value, state.solverState{p}] = ...
                        params.solver(net.params(p).value, state.solverState{p}, ...
                        parDer, params.solverOpts, thisLR) ;
                end
            end
        otherwise
            error('Unknown training method ''%s'' for parameter ''%s''.', ...
                net.params(p).trainMethod, ...
                net.params(p).name) ;
    end
end

% -------------------------------------------------------------------------
function stats = accumulateStats(stats_)
% -------------------------------------------------------------------------

for s = {'train', 'val'}
    s = char(s) ;
    total = 0 ;
    
    % initialize stats stucture with same fields and same order as
    % stats_{1}
    stats__ = stats_{1} ;
    names = fieldnames(stats__.(s))' ;
    values = zeros(1, numel(names)) ;
    fields = cat(1, names, num2cell(values)) ;
    stats.(s) = struct(fields{:}) ;
    
    for g = 1:numel(stats_)
        stats__ = stats_{g} ;
        num__ = stats__.(s).num ;
        total = total + num__ ;
        
        for f = setdiff(fieldnames(stats__.(s))', 'num')
            f = char(f) ;
            stats.(s).(f) = stats.(s).(f) + stats__.(s).(f) * num__ ;
            
            if g == numel(stats_)
                stats.(s).(f) = stats.(s).(f) / total ;
            end
        end
    end
    stats.(s).num = total ;
end

% -------------------------------------------------------------------------
function stats = extractStats(stats, net)
% -------------------------------------------------------------------------
sel = find(cellfun(@(x) isa(x,'dagnn.Loss'), {net.layers.block})) ;
for i = 1:numel(sel)
    if net.layers(sel(i)).block.ignoreAverage, continue; end;
    stats.(net.layers(sel(i)).outputs{1}) = net.layers(sel(i)).block.average ;
end

% -------------------------------------------------------------------------
function saveState(fileName, net_, state)
% -------------------------------------------------------------------------
net = net_.saveobj() ;
fprintf('Save %s\n', fileName);
save(fileName, 'net', 'state') ;

% -------------------------------------------------------------------------
function saveStats(fileName, stats)
% -------------------------------------------------------------------------
fprintf('Save %s\n', fileName);
if exist(fileName)
    save(fileName, 'stats', '-append') ;
else
    save(fileName, 'stats') ;
end

% -------------------------------------------------------------------------
function [net, state, stats] = loadState(fileName)
% -------------------------------------------------------------------------
load(fileName, 'net', 'state', 'stats') ;
net = dagnn.DagNN.loadobj(net) ;
if isempty(whos('stats'))
    error('Epoch ''%s'' was only partially saved. Delete this file and try again.', ...
        fileName) ;
end

% -------------------------------------------------------------------------
function epoch = findLastCheckpoint(modelDir)
% -------------------------------------------------------------------------
list = dir(fullfile(modelDir, 'net-epoch-*.mat')) ;
tokens = regexp({list.name}, 'net-epoch-([\d]+).mat', 'tokens') ;
epoch = cellfun(@(x) sscanf(x{1}{1}, '%d'), tokens) ;
epoch = max([epoch 0]) ;

% -------------------------------------------------------------------------
function switchFigure(n)
% -------------------------------------------------------------------------
if get(0,'CurrentFigure') ~= n
    try
        set(0,'CurrentFigure',n) ;
    catch
        figure(n) ;
    end
end

% -------------------------------------------------------------------------
function clearMex()
% -------------------------------------------------------------------------
clear vl_tmove vl_imreadjpeg ;

% -------------------------------------------------------------------------
function prepareGPUs(opts, cold)
% -------------------------------------------------------------------------
numGpus = numel(opts.gpus) ;
if numGpus > 1
    % check parallel pool integrity as it could have timed out
    pool = gcp('nocreate') ;
    if ~isempty(pool) && pool.NumWorkers ~= numGpus
        delete(pool) ;
    end
    pool = gcp('nocreate') ;
    if isempty(pool)
        parpool('local', numGpus) ;
        cold = true ;
    end
    
end
if numGpus >= 1 && cold
    fprintf('%s: resetting GPU\n', mfilename)
    clearMex() ;
    if numGpus == 1
        gpuDevice(opts.gpus)
    else
        spmd
            clearMex() ;
            gpuDevice(opts.gpus(labindex))
        end
    end
end

function fn = getBatch
% bopts = struct('numGpus', numel(opts.train.gpus)) ;
fn = @(opts,x,y) getDagNNBatch(opts,x,y) ;

function inputs = getDagNNBatch(opts, imdb, batch)
input = imdb.inputs_shadow(:,:,:,batch);
% ratio = imdb.inputs_ratio(:,:,:,batch);
% ratio_down2 = imdb.inputs_ratio2(:,:,:,batch);
% ratio_down4 = imdb.inputs_ratio4(:,:,:,batch);
% ratio_down8 = imdb.inputs_ratio8(:,:,:,batch);
label = imdb.inputs_label(:,:,:,batch);
label_down2 = imdb.inputs_label2(:,:,:,batch);
label_down4 = imdb.inputs_label4(:,:,:,batch);
label_down8 = imdb.inputs_label8(:,:,:,batch);

% convert to GPU array
if opts.gpus > 0
    input = gpuArray(input) ;
%     ratio = gpuArray(ratio) ;
%     ratio_down2 = gpuArray(ratio_down2) ;
%     ratio_down4 = gpuArray(ratio_down4) ;
%     ratio_down8 = gpuArray(ratio_down8) ;
    label = gpuArray(label) ;
    label_down2 = gpuArray(label_down2) ;
    label_down4 = gpuArray(label_down4) ;
    label_down8 = gpuArray(label_down8) ;    
end

% dagnn input
% inputs = {'input',input, 'labels',label};
% inputs = {'input',input, 'ratio',ratio, 'ratio_down2',ratio_down2,'ratio_down4',ratio_down4,'ratio_down8',ratio_down8,'labels',label};
% inputs = {'input',input,'ratio_down2',ratio_down2,'ratio_down4',ratio_down4,'ratio_down8',ratio_down8,'labels',label};
% inputs = {'input',input,'ratio_down2',ratio_down2,'ratio_down4',ratio_down4,'ratio_down8',ratio_down8,...
%          'level1_label',label_down8,'level2_label',label_down4,'level3_label',label_down2,'level4_label',label};
inputs = {'input',input, 'level1_label',label_down8,'level2_label',label_down4,'level3_label',label_down2,'level4_label',label};