package org.deeplearning4j.spark.impl.vanilla;

import lombok.AllArgsConstructor;
import org.apache.spark.broadcast.Broadcast;
import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;
import org.deeplearning4j.spark.api.TrainingWorker;
import org.deeplearning4j.spark.api.WorkerConfiguration;
import org.deeplearning4j.spark.api.worker.NetBroadcastTuple;
import org.nd4j.linalg.dataset.api.DataSet;

/**
 * Created by Alex on 14/06/2016.
 */
@AllArgsConstructor
public class VanillaTrainingWorker implements TrainingWorker<VanillaTrainingResult> {

    private final Broadcast<NetBroadcastTuple> broadcast;
    private final boolean saveUpdater;
    private final WorkerConfiguration configuration;

    @Override
    public MultiLayerNetwork getInitialModel() {
        NetBroadcastTuple tuple = broadcast.getValue();
        MultiLayerNetwork net = new MultiLayerNetwork(tuple.getConfiguration());
        net.init(); //TODO: Replace this with net.init(INDArray params) once that is implemented

        //Can't have shared parameter array across executors for parameter averaging, hence the .dup()
        net.setParameters(tuple.getParameters().dup());

        if(tuple.getUpdater() != null){
            net.setUpdater(tuple.getUpdater());
        }

        return net;
    }

    @Override
    public VanillaTrainingResult processMinibatch(DataSet dataSet, MultiLayerNetwork network, boolean isLast) {
        network.fit(dataSet);

        if(isLast) return getFinalResult(network);

        return null;
    }

    @Override
    public VanillaTrainingResult getFinalResult(MultiLayerNetwork network) {
        //TODO: don't want to use java serialization for updater, in case worker is using cuda and master is using native, etc
        return new VanillaTrainingResult(network.params(), (saveUpdater ? network.getUpdater() : null), network.score());
    }

    @Override
    public WorkerConfiguration getDataConfiguration() {
        return configuration;
    }
}

@@ -899,106 +899,8 @@
 ));%0A
-        net.init(); //TODO: Replace this with net.init(INDArray params) once that is implemented%0A%0A
     
@@ -995,40 +995,62 @@
 the 
-.dup()%0A        net.setParameters
+'true' for clone parameters array arg%0A        net.init
 (tup
@@ -1071,14 +1071,14 @@
 rs()
-.dup()
+, true
 );%0A%0A
@@ -1162,18 +1162,62 @@
 pdater()
-);
+.clone()); //Again: can'h have shared updaters
 %0A       

Switch Vanilla training worker to use new init(INDArray,boolean)