cavis/cavis-dnn/cavis-dnn-core/src/test/java/org/deeplearning4j/optimize/solver/BackTrackLineSearchTest.java

/*
 *  ******************************************************************************
 *  *
 *  *
 *  * This program and the accompanying materials are made available under the
 *  * terms of the Apache License, Version 2.0 which is available at
 *  * https://www.apache.org/licenses/LICENSE-2.0.
 *  *
 *  *  See the NOTICE file distributed with this work for additional
 *  *  information regarding copyright ownership.
 *  * Unless required by applicable law or agreed to in writing, software
 *  * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
 *  * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
 *  * License for the specific language governing permissions and limitations
 *  * under the License.
 *  *
 *  * SPDX-License-Identifier: Apache-2.0
 *  *****************************************************************************
 */

package org.deeplearning4j.optimize.solver;

import lombok.val;
import org.deeplearning4j.BaseDL4JTest;
import org.deeplearning4j.datasets.iterator.impl.IrisDataSetIterator;
import org.deeplearning4j.nn.api.OptimizationAlgorithm;
import org.deeplearning4j.nn.conf.NeuralNetConfiguration;
import org.deeplearning4j.nn.conf.layers.DenseLayer;
import org.deeplearning4j.nn.layers.OutputLayer;
import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;
import org.deeplearning4j.nn.weights.WeightInit;
import org.deeplearning4j.nn.workspace.LayerWorkspaceMgr;
import org.deeplearning4j.optimize.api.TrainingListener;
import org.deeplearning4j.optimize.listeners.ScoreIterationListener;
import org.deeplearning4j.optimize.solvers.BackTrackLineSearch;
import org.deeplearning4j.optimize.stepfunctions.DefaultStepFunction;
import org.deeplearning4j.optimize.stepfunctions.NegativeDefaultStepFunction;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.nd4j.linalg.activations.Activation;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.dataset.DataSet;
import org.nd4j.linalg.dataset.api.iterator.DataSetIterator;
import org.nd4j.linalg.factory.Nd4j;
import org.nd4j.linalg.learning.config.Adam;
import org.nd4j.linalg.lossfunctions.LossFunctions;

import java.util.Collections;

import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertTrue;

/**
 * @author Adam Gibson
 */
public class BackTrackLineSearchTest extends BaseDL4JTest {
    private DataSetIterator irisIter;
    private DataSet irisData;

    @BeforeEach
    public void before() {
        if (irisIter == null) {
            irisIter = new IrisDataSetIterator(5, 5);
        }
        if (irisData == null) {
            irisData = irisIter.next();
            irisData.normalizeZeroMeanZeroUnitVariance();
        }
    }



    @Test
    public void testSingleMinLineSearch() throws Exception {
        OutputLayer layer = getIrisLogisticLayerConfig(Activation.SOFTMAX, 100,
                        LossFunctions.LossFunction.NEGATIVELOGLIKELIHOOD);
        int nParams = (int)layer.numParams();
        layer.setBackpropGradientsViewArray(Nd4j.create(1, nParams));
        layer.setInput(irisData.getFeatures(), LayerWorkspaceMgr.noWorkspaces());
        layer.setLabels(irisData.getLabels());
        layer.computeGradientAndScore(LayerWorkspaceMgr.noWorkspaces());

        BackTrackLineSearch lineSearch = new BackTrackLineSearch(layer, layer.getOptimizer());
        double step = lineSearch.optimize(layer.getModelParams(), layer.gradient().gradient(), layer.gradient().gradient(), LayerWorkspaceMgr.noWorkspacesImmutable());

        assertEquals(1.0, step, 1e-3);
    }

    @Test
    public void testSingleMaxLineSearch() throws Exception {
        double score1, score2;

        OutputLayer layer = getIrisLogisticLayerConfig(Activation.SOFTMAX, 100,
                        LossFunctions.LossFunction.NEGATIVELOGLIKELIHOOD);
        int nParams = (int)layer.numParams();
        layer.setBackpropGradientsViewArray(Nd4j.create(1, nParams));
        layer.setInput(irisData.getFeatures(), LayerWorkspaceMgr.noWorkspaces());
        layer.setLabels(irisData.getLabels());
        layer.computeGradientAndScore(LayerWorkspaceMgr.noWorkspaces());
        score1 = layer.getScore();

        BackTrackLineSearch lineSearch =
                        new BackTrackLineSearch(layer, new NegativeDefaultStepFunction(), layer.getOptimizer());
        double step = lineSearch.optimize(layer.getModelParams(), layer.gradient().gradient(), layer.gradient().gradient(), LayerWorkspaceMgr.noWorkspacesImmutable());

        assertEquals(1.0, step, 1e-3);
    }


    @Test
    public void testMultMinLineSearch() throws Exception {
        double score1, score2;

        OutputLayer layer = getIrisLogisticLayerConfig(Activation.SOFTMAX, 100,
                        LossFunctions.LossFunction.NEGATIVELOGLIKELIHOOD);
        int nParams = (int)layer.numParams();
        layer.setBackpropGradientsViewArray(Nd4j.create(1, nParams));
        layer.setInput(irisData.getFeatures(), LayerWorkspaceMgr.noWorkspaces());
        layer.setLabels(irisData.getLabels());
        layer.computeGradientAndScore(LayerWorkspaceMgr.noWorkspaces());
        score1 = layer.getScore();
        INDArray origGradient = layer.gradient().gradient().dup();

        NegativeDefaultStepFunction sf = new NegativeDefaultStepFunction();
        BackTrackLineSearch lineSearch = new BackTrackLineSearch(layer, sf, layer.getOptimizer());
        double step = lineSearch.optimize(layer.getModelParams(), layer.gradient().gradient(), layer.gradient().gradient(), LayerWorkspaceMgr.noWorkspacesImmutable());
        INDArray currParams = layer.getModelParams();
        sf.step(currParams, origGradient, step);
        layer.setParamsTable(currParams);
        layer.computeGradientAndScore(LayerWorkspaceMgr.noWorkspaces());

        score2 = layer.getScore();

        assertTrue(score1 > score2, "score1=" + score1 + ", score2=" + score2);

    }

    @Test
    public void testMultMaxLineSearch() throws Exception {
        double score1, score2;

        irisData.normalizeZeroMeanZeroUnitVariance();
        OutputLayer layer = getIrisLogisticLayerConfig(Activation.SOFTMAX, 100, LossFunctions.LossFunction.MCXENT);
        int nParams = (int)layer.numParams();
        layer.setBackpropGradientsViewArray(Nd4j.create(1, nParams));
        layer.setInput(irisData.getFeatures(), LayerWorkspaceMgr.noWorkspaces());
        layer.setLabels(irisData.getLabels());
        layer.computeGradientAndScore(LayerWorkspaceMgr.noWorkspaces());
        score1 = layer.getScore();
        INDArray origGradient = layer.gradient().gradient().dup();

        DefaultStepFunction sf = new DefaultStepFunction();
        BackTrackLineSearch lineSearch = new BackTrackLineSearch(layer, sf, layer.getOptimizer());
        double step = lineSearch.optimize(layer.getModelParams().dup(), layer.gradient().gradient().dup(),
                        layer.gradient().gradient().dup(), LayerWorkspaceMgr.noWorkspacesImmutable());

        INDArray currParams = layer.getModelParams();
        sf.step(currParams, origGradient, step);
        layer.setParamsTable(currParams);
        layer.computeGradientAndScore(LayerWorkspaceMgr.noWorkspaces());
        score2 = layer.getScore();

        assertTrue(score1 < score2, "score1 = " + score1 + ", score2 = " + score2);
    }

    private static OutputLayer getIrisLogisticLayerConfig(Activation activationFunction, int maxIterations,
                    LossFunctions.LossFunction lossFunction) {
        NeuralNetConfiguration conf =
                        NeuralNetConfiguration.builder().seed(12345L).miniBatch(true)
                                        .maxNumLineSearchIterations(maxIterations)
                                        .layer(org.deeplearning4j.nn.conf.layers.OutputLayer.builder().lossFunction(lossFunction)
                                                        .nIn(4).nOut(3).activation(activationFunction)
                                                        .weightInit(WeightInit.XAVIER).build())
                                        .build();

        val numParams = conf.getFirstLayer().initializer().numParams(conf);
        INDArray params = Nd4j.create(1, numParams);
        return (OutputLayer) conf.getFirstLayer().instantiate(conf, null, 0, params, true, params.dataType());
    }

    ///////////////////////////////////////////////////////////////////////////

    @Test
    public void testBackTrackLineGradientDescent() {
        OptimizationAlgorithm optimizer = OptimizationAlgorithm.LINE_GRADIENT_DESCENT;

        DataSetIterator irisIter = new IrisDataSetIterator(1, 1);
        DataSet data = irisIter.next();

        MultiLayerNetwork network = new MultiLayerNetwork(getIrisMultiLayerConfig(Activation.SIGMOID, optimizer));
        network.init();
        TrainingListener listener = new ScoreIterationListener(10);
        network.addTrainingListeners(Collections.singletonList(listener));
        double oldScore = network.score(data);
        for( int i=0; i<100; i++ ) {
            network.fit(data.getFeatures(), data.getLabels());
        }
        double score = network.getScore();
        assertTrue(score < oldScore);
    }

    @Test
    public void testBackTrackLineCG() {
        OptimizationAlgorithm optimizer = OptimizationAlgorithm.CONJUGATE_GRADIENT;

        DataSet data = irisIter.next();
        data.normalizeZeroMeanZeroUnitVariance();
        MultiLayerNetwork network = new MultiLayerNetwork(getIrisMultiLayerConfig(Activation.RELU, optimizer));
        network.init();
        TrainingListener listener = new ScoreIterationListener(10);
        network.addTrainingListeners(Collections.singletonList(listener));
        double firstScore = network.score(data);

        for( int i=0; i<5; i++ ) {
            network.fit(data.getFeatures(), data.getLabels());
        }
        double score = network.getScore();
        assertTrue(score < firstScore);

    }

    @Test
    public void testBackTrackLineLBFGS() {
        OptimizationAlgorithm optimizer = OptimizationAlgorithm.LBFGS;
        DataSet data = irisIter.next();
        data.normalizeZeroMeanZeroUnitVariance();
        MultiLayerNetwork network = new MultiLayerNetwork(getIrisMultiLayerConfig(Activation.RELU, optimizer));
        network.init();
        TrainingListener listener = new ScoreIterationListener(10);
        network.addTrainingListeners(Collections.singletonList(listener));
        double oldScore = network.score(data);

        for( int i=0; i<5; i++ ) {
            network.fit(data.getFeatures(), data.getLabels());
        }
        double score = network.getScore();
        assertTrue(score < oldScore);

    }

    private static NeuralNetConfiguration getIrisMultiLayerConfig(Activation activationFunction, OptimizationAlgorithm optimizer) {
        NeuralNetConfiguration conf = NeuralNetConfiguration.builder().optimizationAlgo(optimizer)
                        .updater(new Adam(0.01)).seed(12345L).list()
                        .layer(0, DenseLayer.builder().nIn(4).nOut(100).weightInit(WeightInit.XAVIER)
                                        .activation(activationFunction).build())
                        .layer(1, org.deeplearning4j.nn.conf.layers.OutputLayer.builder().lossFunction(
                                        LossFunctions.LossFunction.MCXENT).nIn(100).nOut(3)
                                                        .weightInit(WeightInit.XAVIER).activation(Activation.SOFTMAX)
                                                        .build())
                        .build();


        return conf;
    }

}