org.deeplearning4j.spark.impl.multilayer.IterativeReduceFlatMap Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of dl4j-spark_2.11 Show documentation
There is a newer version: 1.0.0-beta_spark_2
/*
 *
 *  * Copyright 2015 Skymind,Inc.
 *  *
 *  *    Licensed under the Apache License, Version 2.0 (the "License");
 *  *    you may not use this file except in compliance with the License.
 *  *    You may obtain a copy of the License at
 *  *
 *  *        http://www.apache.org/licenses/LICENSE-2.0
 *  *
 *  *    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.
 *
 */

package org.deeplearning4j.spark.impl.multilayer;

import org.apache.spark.Accumulator;
import org.apache.spark.api.java.function.FlatMapFunction;
import org.apache.spark.broadcast.Broadcast;
import org.deeplearning4j.nn.api.Updater;
import org.deeplearning4j.nn.conf.MultiLayerConfiguration;
import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;
import org.deeplearning4j.optimize.listeners.ScoreIterationListener;
import org.deeplearning4j.spark.impl.common.BestScoreIterationListener;
import org.deeplearning4j.spark.impl.common.misc.ScoreReport;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.dataset.DataSet;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import scala.Tuple3;

import java.util.ArrayList;
import java.util.Collections;
import java.util.Iterator;
import java.util.List;

/**
 * Iterative reduce with
 * flat map using map partitions
 *
 * @author Adam Gibson
 */


public class IterativeReduceFlatMap implements FlatMapFunction,Tuple3> {
    protected static Logger log = LoggerFactory.getLogger(IterativeReduceFlatMap.class);

    protected String json;
    protected Broadcast params;
    protected Broadcast updater;
    protected final Accumulator best_score_acc;

    /**
     * Pass in network configuration as json, broadcast parameters, broadcast updater and the bestScore to
     * fit mapped data set and update and return parameters, updater and best score.
     * @param json newtork string configuration
     * @param params broadcasted parameters to reload into network
     * @param updater broadcasted updaters to reload into network
     * @param bestScoreAcc accumulator which tracks best score seen
     */
    public IterativeReduceFlatMap(String json, Broadcast params, Broadcast updater,
                                  Accumulator bestScoreAcc) {
        this.json = json;
        this.params = params;
        this.updater = updater;
        if(updater.getValue() == null)
            throw new IllegalArgumentException("Updater shouldn't be null");
        this.best_score_acc = bestScoreAcc;
    }

    @Override
    public Iterable> call(Iterator dataSetIterator) throws Exception {
        if (!dataSetIterator.hasNext()) {
            return Collections.emptyList();
        }
        List collect = new ArrayList<>();
        while (dataSetIterator.hasNext()) {
            collect.add(dataSetIterator.next());
        }

        DataSet data = DataSet.merge(collect, false);
        if (log.isDebugEnabled()) {
            log.debug("Training on {} examples with data {}", data.numExamples(), data.labelCounts());
        }

        //Need to clone: parameters and updaters are mutable values -> .getValue() object will be shared by ALL executors on the same machine!
        INDArray val = params.getValue().dup();
        Updater upd = updater.getValue().clone();

        MultiLayerNetwork network = new MultiLayerNetwork(MultiLayerConfiguration.fromJson(json));
        network.setInitDone(true);
        network.init();
        network.setListeners(new ScoreIterationListener(1), new BestScoreIterationListener(best_score_acc));
        if (val.length() != network.numParams(false))
            throw new IllegalStateException("Network did not have same number of parameters as the broadcasted set parameters");
        network.setParameters(val);
        network.setUpdater(upd);
        network.fit(data);
        ScoreReport report = new ScoreReport();
        report.setS(network.score());
        report.setM(Runtime.getRuntime().maxMemory());
        return Collections.singletonList(new Tuple3<>(network.params(false), network.getUpdater(), report));
    }
}