chore: import upstream snapshot with attribution

deeplearning4j/deeplearning4j-nn/pom.xml

@@ -0,0 +1,117 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!--
+  ~ /* ******************************************************************************
+  ~  *
+  ~  *
+  ~  * 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
+  ~  ******************************************************************************/
+  -->
+
+<project xmlns="http://maven.apache.org/POM/4.0.0"
+    xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+    xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
+
+    <modelVersion>4.0.0</modelVersion>
+
+    <parent>
+        <groupId>org.eclipse.deeplearning4j</groupId>
+        <artifactId>deeplearning4j-parent</artifactId>
+        <version>1.0.0-SNAPSHOT</version>
+    </parent>
+
+    <artifactId>deeplearning4j-nn</artifactId>
+
+    <name>deeplearning4j-nn</name>
+
+    <build>
+        <plugins>
+            <plugin>
+                <groupId>org.moditect</groupId>
+                <artifactId>moditect-maven-plugin</artifactId>
+            </plugin>
+        </plugins>
+    </build>
+
+    <properties>
+        <module.name>deeplearning4j.nn</module.name>
+    </properties>
+
+    <dependencies>
+        <dependency>
+            <groupId>org.eclipse.deeplearning4j</groupId>
+            <artifactId>deeplearning4j-utility-iterators</artifactId>
+            <version>${project.version}</version>
+        </dependency>
+
+        <dependency>
+            <groupId>commons-io</groupId>
+            <artifactId>commons-io</artifactId>
+            <version>${commonsio.version}</version>
+        </dependency>
+        <!-- ND4J API -->
+        <dependency>
+            <groupId>org.eclipse.deeplearning4j</groupId>
+            <artifactId>nd4j-api</artifactId>
+            <version>${nd4j.version}</version>
+        </dependency>
+        <dependency>
+            <groupId>org.eclipse.deeplearning4j</groupId>
+            <artifactId>nd4j-native-api</artifactId>
+            <version>${nd4j.version}</version>
+        </dependency>
+        <dependency>
+            <groupId>org.eclipse.deeplearning4j</groupId>
+            <artifactId>nd4j-common</artifactId>
+            <version>${nd4j.version}</version>
+        </dependency>
+        <dependency>
+            <groupId>com.google.code.gson</groupId>
+            <artifactId>gson</artifactId>
+            <version>${gson.version}</version>
+        </dependency>
+        <!-- ND4J Shaded Jackson Dependency -->
+        <dependency>
+            <groupId>org.eclipse.deeplearning4j</groupId>
+            <artifactId>jackson</artifactId>
+            <version>${nd4j.version}</version>
+        </dependency>
+        <!-- oshi: Used for collecting system information for memory crash dump reporting -->
+        <dependency>
+            <groupId>com.github.oshi</groupId>
+            <artifactId>oshi-core</artifactId>
+            <version>${oshi.version}</version>
+        </dependency>
+        <dependency>
+            <groupId>ch.qos.logback</groupId>
+            <artifactId>logback-classic</artifactId>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>it.unimi.dsi</groupId>
+            <artifactId>fastutil</artifactId>
+            <version>${fastutil.version}</version>
+        </dependency>
+
+
+        <dependency>
+            <groupId>org.eclipse.deeplearning4j</groupId>
+            <artifactId>resources</artifactId>
+            <version>${project.version}</version>
+        </dependency>
+
+    </dependencies>
+
+
+</project>

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/earlystopping/EarlyStoppingConfiguration.java

@@ -0,0 +1,167 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.earlystopping;
+
+import lombok.Data;
+import lombok.NoArgsConstructor;
+import org.deeplearning4j.earlystopping.saver.InMemoryModelSaver;
+import org.deeplearning4j.earlystopping.scorecalc.ScoreCalculator;
+import org.deeplearning4j.earlystopping.termination.EpochTerminationCondition;
+import org.deeplearning4j.earlystopping.termination.IterationTerminationCondition;
+import org.deeplearning4j.exception.DL4JInvalidConfigException;
+import org.deeplearning4j.nn.api.Model;
+import org.nd4j.common.function.Supplier;
+
+import java.io.Serializable;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+
+@Data
+@NoArgsConstructor
+public class EarlyStoppingConfiguration<T extends Model> implements Serializable {
+
+    private EarlyStoppingModelSaver<T> modelSaver;
+    private List<EpochTerminationCondition> epochTerminationConditions;
+    private List<IterationTerminationCondition> iterationTerminationConditions;
+    private boolean saveLastModel;
+    private int evaluateEveryNEpochs;
+    private ScoreCalculator<T> scoreCalculator;
+    private Supplier<ScoreCalculator> scoreCalculatorSupplier;
+
+    private EarlyStoppingConfiguration(Builder<T> builder) {
+        this.modelSaver = builder.modelSaver;
+        this.epochTerminationConditions = builder.epochTerminationConditions;
+        this.iterationTerminationConditions = builder.iterationTerminationConditions;
+        this.saveLastModel = builder.saveLastModel;
+        this.evaluateEveryNEpochs = builder.evaluateEveryNEpochs;
+        this.scoreCalculator = builder.scoreCalculator;
+        this.scoreCalculatorSupplier = builder.scoreCalculatorSupplier;
+    }
+
+    public ScoreCalculator<T> getScoreCalculator(){
+        if(scoreCalculatorSupplier != null){
+            return scoreCalculatorSupplier.get();
+        }
+        return scoreCalculator;
+    }
+
+
+    public void validate() {
+        if(scoreCalculator == null && scoreCalculatorSupplier == null) {
+            throw new DL4JInvalidConfigException("A score calculator or score calculator supplier must be defined.");
+        }
+
+        if(modelSaver == null) {
+            throw new DL4JInvalidConfigException("A model saver must be defined");
+        }
+
+        boolean hasTermination = false;
+        if(iterationTerminationConditions != null && !iterationTerminationConditions.isEmpty()) {
+            hasTermination = true;
+        }
+
+        else if(epochTerminationConditions != null && !epochTerminationConditions.isEmpty()) {
+            hasTermination = true;
+        }
+
+        if(!hasTermination) {
+            throw new DL4JInvalidConfigException("No termination conditions defined.");
+        }
+    }
+
+
+    public static class Builder<T extends Model> {
+
+        private EarlyStoppingModelSaver<T> modelSaver = new InMemoryModelSaver<>();
+        private List<EpochTerminationCondition> epochTerminationConditions = new ArrayList<>();
+        private List<IterationTerminationCondition> iterationTerminationConditions = new ArrayList<>();
+        private boolean saveLastModel = false;
+        private int evaluateEveryNEpochs = 1;
+        private ScoreCalculator<T> scoreCalculator;
+        private Supplier<ScoreCalculator> scoreCalculatorSupplier;
+
+
+        /** How should models be saved? (Default: in memory)*/
+        public Builder<T> modelSaver(EarlyStoppingModelSaver<T> modelSaver) {
+            this.modelSaver = modelSaver;
+            return this;
+        }
+
+        /** Termination conditions to be evaluated every N epochs, with N set by evaluateEveryNEpochs option */
+        public Builder<T> epochTerminationConditions(EpochTerminationCondition... terminationConditions) {
+            epochTerminationConditions.clear();
+            Collections.addAll(epochTerminationConditions, terminationConditions);
+            return this;
+        }
+
+        /** Termination conditions to be evaluated every N epochs, with N set by evaluateEveryNEpochs option */
+        public Builder<T> epochTerminationConditions(List<EpochTerminationCondition> terminationConditions) {
+            this.epochTerminationConditions = terminationConditions;
+            return this;
+        }
+
+        /** Termination conditions to be evaluated every iteration (minibatch)*/
+        public Builder<T> iterationTerminationConditions(IterationTerminationCondition... terminationConditions) {
+            iterationTerminationConditions.clear();
+            Collections.addAll(iterationTerminationConditions, terminationConditions);
+            return this;
+        }
+
+        /** Save the last model? If true: save the most recent model at each epoch, in addition to the best
+         * model (whenever the best model improves). If false: only save the best model. Default: false
+         * Useful for example if you might want to continue training after a max-time terminatino condition
+         * occurs.
+         */
+        public Builder<T> saveLastModel(boolean saveLastModel) {
+            this.saveLastModel = saveLastModel;
+            return this;
+        }
+
+        /** How frequently should evaluations be conducted (in terms of epochs)? Defaults to every (1) epochs. */
+        public Builder<T> evaluateEveryNEpochs(int everyNEpochs) {
+            this.evaluateEveryNEpochs = everyNEpochs;
+            return this;
+        }
+
+        /** Score calculator. Used to calculate a score (such as loss function on a test set), every N epochs,
+         * where N is set by {@link #evaluateEveryNEpochs}
+         */
+        public Builder<T> scoreCalculator(ScoreCalculator scoreCalculator) {
+            this.scoreCalculator = scoreCalculator;
+            return this;
+        }
+
+        /** Score calculator. Used to calculate a score (such as loss function on a test set), every N epochs,
+         * where N is set by {@link #evaluateEveryNEpochs}
+         */
+        public Builder<T> scoreCalculator(Supplier<ScoreCalculator> scoreCalculatorSupplier){
+            this.scoreCalculatorSupplier = scoreCalculatorSupplier;
+            return this;
+        }
+
+        /** Create the early stopping configuration */
+        public EarlyStoppingConfiguration<T> build() {
+            return new EarlyStoppingConfiguration<>(this);
+        }
+
+    }
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/earlystopping/EarlyStoppingModelSaver.java

@@ -0,0 +1,55 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.earlystopping;
+
+import org.deeplearning4j.earlystopping.saver.InMemoryModelSaver;
+import org.deeplearning4j.earlystopping.saver.LocalFileGraphSaver;
+import org.deeplearning4j.earlystopping.saver.LocalFileModelSaver;
+import org.deeplearning4j.nn.api.Model;
+import org.nd4j.shade.jackson.annotation.JsonInclude;
+import org.nd4j.shade.jackson.annotation.JsonSubTypes;
+import org.nd4j.shade.jackson.annotation.JsonTypeInfo;
+
+import java.io.IOException;
+import java.io.Serializable;
+
+@JsonInclude(JsonInclude.Include.NON_NULL)
+@JsonSubTypes(value = {@JsonSubTypes.Type(value = InMemoryModelSaver.class, name = "InMemoryModelSaver"),
+                @JsonSubTypes.Type(value = LocalFileGraphSaver.class, name = "LocalFileGraphSaver"),
+                @JsonSubTypes.Type(value = LocalFileModelSaver.class, name = "LocalFileModelSaver"),
+
+})
+@JsonTypeInfo(use = JsonTypeInfo.Id.CLASS, include = JsonTypeInfo.As.PROPERTY, property = "@class")
+public interface EarlyStoppingModelSaver<T extends Model> extends Serializable {
+
+    /** Save the best model (so far) learned during early stopping training */
+    void saveBestModel(T net, double score) throws IOException;
+
+    /** Save the latest (most recent) model learned during early stopping */
+    void saveLatestModel(T net, double score) throws IOException;
+
+    /** Retrieve the best model that was previously saved */
+    T getBestModel() throws IOException;
+
+    /** Retrieve the most recent model that was previously saved */
+    T getLatestModel() throws IOException;
+
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/earlystopping/EarlyStoppingResult.java

@@ -0,0 +1,67 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.earlystopping;
+
+import lombok.Data;
+import org.deeplearning4j.nn.api.Model;
+
+import java.io.Serializable;
+import java.util.Map;
+
+@Data
+public class EarlyStoppingResult<T extends Model> implements Serializable {
+    public enum TerminationReason {
+        Error, IterationTerminationCondition, EpochTerminationCondition
+    }
+
+    private TerminationReason terminationReason;
+    private String terminationDetails;
+    private Map<Integer, Double> scoreVsEpoch;
+    private int bestModelEpoch;
+    private double bestModelScore;
+    private int totalEpochs;
+    private T bestModel;
+
+    public EarlyStoppingResult(TerminationReason terminationReason, String terminationDetails,
+                    Map<Integer, Double> scoreVsEpoch, int bestModelEpoch, double bestModelScore, int totalEpochs,
+                    T bestModel) {
+        this.terminationReason = terminationReason;
+        this.terminationDetails = terminationDetails;
+        this.scoreVsEpoch = scoreVsEpoch;
+        this.bestModelEpoch = bestModelEpoch;
+        this.bestModelScore = bestModelScore;
+        this.totalEpochs = totalEpochs;
+        this.bestModel = bestModel;
+    }
+
+    @Override
+    public String toString() {
+        return "EarlyStoppingResult(terminationReason=" + terminationReason + ",details=" + terminationDetails
+                        + ",bestModelEpoch=" + bestModelEpoch + ",bestModelScore=" + bestModelScore + ",totalEpochs="
+                        + totalEpochs + ")";
+
+    }
+
+    public T getBestModel() {
+        return bestModel;
+    }
+
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/earlystopping/listener/EarlyStoppingListener.java

@@ -0,0 +1,46 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.earlystopping.listener;
+
+import org.deeplearning4j.earlystopping.EarlyStoppingConfiguration;
+import org.deeplearning4j.earlystopping.EarlyStoppingResult;
+import org.deeplearning4j.nn.api.Model;
+
+public interface EarlyStoppingListener<T extends Model> {
+
+    /**Method to be called when early stopping training is first started
+     */
+    void onStart(EarlyStoppingConfiguration<T> esConfig, T net);
+
+    /**Method that is called at the end of each epoch completed during early stopping training
+     * @param epochNum The number of the epoch just completed (starting at 0)
+     * @param score The score calculated
+     * @param esConfig Configuration
+     * @param net Network (current)
+     */
+    void onEpoch(int epochNum, double score, EarlyStoppingConfiguration<T> esConfig, T net);
+
+    /**Method that is called at the end of early stopping training
+     * @param esResult The early stopping result. Provides details of why early stopping training was terminated, etc
+     */
+    void onCompletion(EarlyStoppingResult<T> esResult);
+
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/earlystopping/saver/InMemoryModelSaver.java

@@ -0,0 +1,69 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.earlystopping.saver;
+
+import org.deeplearning4j.earlystopping.EarlyStoppingModelSaver;
+import org.deeplearning4j.nn.api.Model;
+
+import java.io.IOException;
+
+public class InMemoryModelSaver<T extends Model> implements EarlyStoppingModelSaver<T> {
+
+    private transient T bestModel;
+    private transient T latestModel;
+
+    @Override
+    @SuppressWarnings("unchecked")
+    public void saveBestModel(T net, double score) throws IOException {
+        try {
+            //Necessary because close is protected :S
+            bestModel = (T) (net.getClass().getDeclaredMethod("clone")).invoke(net);
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    @Override
+    @SuppressWarnings("unchecked")
+    public void saveLatestModel(T net, double score) throws IOException {
+        try {
+            //Necessary because close is protected :S
+            latestModel = (T) (net.getClass().getDeclaredMethod("clone")).invoke(net);
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    @Override
+    public T getBestModel() throws IOException {
+        return bestModel;
+    }
+
+    @Override
+    public T getLatestModel() throws IOException {
+        return latestModel;
+    }
+
+    @Override
+    public String toString() {
+        return "InMemoryModelSaver()";
+    }
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/earlystopping/saver/LocalFileGraphSaver.java

@@ -0,0 +1,98 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.earlystopping.saver;
+
+import org.apache.commons.io.FilenameUtils;
+import org.deeplearning4j.earlystopping.EarlyStoppingModelSaver;
+import org.deeplearning4j.nn.graph.ComputationGraph;
+import org.deeplearning4j.util.ModelSerializer;
+
+import java.io.File;
+import java.io.IOException;
+import java.nio.charset.Charset;
+
+public class LocalFileGraphSaver implements EarlyStoppingModelSaver<ComputationGraph> {
+
+    private static final String BEST_GRAPH_BIN = "bestGraph.bin";
+    private static final String LATEST_GRAPH_BIN = "latestGraph.bin";
+
+    private String directory;
+    private Charset encoding;
+
+    /**Constructor that uses default character set for configuration (json) encoding
+     * @param directory Directory to save networks
+     */
+    public LocalFileGraphSaver(String directory) {
+        this(directory, Charset.defaultCharset());
+    }
+
+    /**
+     * @param directory Directory to save networks
+     * @param encoding Character encoding for configuration (json)
+     */
+    public LocalFileGraphSaver(String directory, Charset encoding) {
+        this.directory = directory;
+        this.encoding = encoding;
+
+        File dir = new File(directory);
+        if (!dir.exists()) {
+            dir.mkdirs();
+        }
+    }
+
+    @Override
+    public void saveBestModel(ComputationGraph net, double score) throws IOException {
+        String confOut = FilenameUtils.concat(directory, BEST_GRAPH_BIN);
+        save(net, confOut);
+    }
+
+    @Override
+    public void saveLatestModel(ComputationGraph net, double score) throws IOException {
+        String confOut = FilenameUtils.concat(directory, LATEST_GRAPH_BIN);
+        save(net, confOut);
+    }
+
+    private void save(ComputationGraph net, String confOut) throws IOException {
+        ModelSerializer.writeModel(net, confOut, true);
+    }
+
+    @Override
+    public ComputationGraph getBestModel() throws IOException {
+        String confOut = FilenameUtils.concat(directory, BEST_GRAPH_BIN);
+        return load(confOut);
+    }
+
+    @Override
+    public ComputationGraph getLatestModel() throws IOException {
+        String confOut = FilenameUtils.concat(directory, LATEST_GRAPH_BIN);
+        return load(confOut);
+    }
+
+    private ComputationGraph load(String confOut) throws IOException {
+        ComputationGraph net = ModelSerializer.restoreComputationGraph(confOut);
+        return net;
+    }
+
+    @Override
+    public String toString() {
+        return "LocalFileGraphSaver(dir=" + directory + ")";
+    }
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/earlystopping/saver/LocalFileModelSaver.java

@@ -0,0 +1,101 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.earlystopping.saver;
+
+import org.apache.commons.io.FilenameUtils;
+import org.deeplearning4j.earlystopping.EarlyStoppingModelSaver;
+import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;
+import org.deeplearning4j.util.ModelSerializer;
+
+import java.io.File;
+import java.io.IOException;
+import java.nio.charset.Charset;
+
+public class LocalFileModelSaver implements EarlyStoppingModelSaver<MultiLayerNetwork> {
+
+    private static final String BEST_MODEL_BIN = "bestModel.bin";
+    private static final String LATEST_MODEL_BIN = "latestModel.bin";
+    private String directory;
+    private Charset encoding;
+
+    public LocalFileModelSaver(File directory){
+        this(directory.getAbsolutePath());
+    }
+
+    /**Constructor that uses default character set for configuration (json) encoding
+     * @param directory Directory to save networks
+     */
+    public LocalFileModelSaver(String directory) {
+        this(directory, Charset.defaultCharset());
+    }
+
+    /**
+     * @param directory Directory to save networks
+     * @param encoding Character encoding for configuration (json)
+     */
+    public LocalFileModelSaver(String directory, Charset encoding) {
+        this.directory = directory;
+        this.encoding = encoding;
+
+        File dir = new File(directory);
+        if (!dir.exists()) {
+            dir.mkdirs();
+        }
+    }
+
+    @Override
+    public void saveBestModel(MultiLayerNetwork net, double score) throws IOException {
+        String confOut = FilenameUtils.concat(directory, BEST_MODEL_BIN);
+        save(net, confOut);
+    }
+
+    @Override
+    public void saveLatestModel(MultiLayerNetwork net, double score) throws IOException {
+        String confOut = FilenameUtils.concat(directory, LATEST_MODEL_BIN);
+        save(net, confOut);
+    }
+
+    @Override
+    public MultiLayerNetwork getBestModel() throws IOException {
+        String confOut = FilenameUtils.concat(directory, BEST_MODEL_BIN);
+        return load(confOut);
+    }
+
+    @Override
+    public MultiLayerNetwork getLatestModel() throws IOException {
+        String confOut = FilenameUtils.concat(directory, LATEST_MODEL_BIN);
+        return load(confOut);
+    }
+
+    private void save(MultiLayerNetwork net, String modelName) throws IOException {
+        ModelSerializer.writeModel(net, modelName, true);
+    }
+
+    private MultiLayerNetwork load(String modelName) throws IOException {
+        MultiLayerNetwork net = ModelSerializer.restoreMultiLayerNetwork(modelName);
+        return net;
+    }
+
+    @Override
+    public String toString() {
+        return "LocalFileModelSaver(dir=" + directory + ")";
+    }
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/earlystopping/scorecalc/AutoencoderScoreCalculator.java

@@ -0,0 +1,99 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.earlystopping.scorecalc;
+
+import org.deeplearning4j.earlystopping.scorecalc.base.BaseScoreCalculator;
+import org.deeplearning4j.nn.api.Layer;
+import org.deeplearning4j.nn.api.Model;
+import org.deeplearning4j.nn.graph.ComputationGraph;
+import org.deeplearning4j.nn.layers.feedforward.autoencoder.AutoEncoder;
+import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;
+import org.nd4j.evaluation.regression.RegressionEvaluation;
+import org.nd4j.evaluation.regression.RegressionEvaluation.Metric;
+import org.nd4j.linalg.api.ndarray.INDArray;
+import org.nd4j.linalg.dataset.api.iterator.DataSetIterator;
+import org.deeplearning4j.nn.workspace.LayerWorkspaceMgr;
+
+public class AutoencoderScoreCalculator extends BaseScoreCalculator<Model> {
+
+    protected final Metric metric;
+    protected RegressionEvaluation evaluation;
+
+    public AutoencoderScoreCalculator(Metric metric, DataSetIterator iterator){
+        super(iterator);
+        this.metric = metric;
+    }
+
+    @Override
+    protected void reset() {
+        evaluation = new RegressionEvaluation();
+    }
+
+    @Override
+    protected INDArray output(Model net, INDArray input, INDArray fMask, INDArray lMask) {
+
+        Layer l;
+        if(net instanceof MultiLayerNetwork) {
+            MultiLayerNetwork network = (MultiLayerNetwork)net;
+            l = network.getLayer(0);
+        } else {
+            ComputationGraph network = (ComputationGraph)net;
+            l = network.getLayer(0);
+        }
+
+        if (!(l instanceof AutoEncoder)) {
+            throw new UnsupportedOperationException("Can only score networks with autoencoder layers as first layer -" +
+                    " got " + l.getClass().getSimpleName());
+        }
+        AutoEncoder ae = (AutoEncoder) l;
+
+        LayerWorkspaceMgr workspaceMgr = LayerWorkspaceMgr.noWorkspaces();
+        INDArray encode = ae.encode(input, false, workspaceMgr);
+        return ae.decode(encode, workspaceMgr);
+    }
+
+    @Override
+    protected INDArray[] output(Model network, INDArray[] input, INDArray[] fMask, INDArray[] lMask) {
+        return new INDArray[]{output(network, get0(input), get0(fMask), get0(lMask))};
+    }
+
+    @Override
+    protected double scoreMinibatch(Model network, INDArray features, INDArray labels, INDArray fMask,
+                                    INDArray lMask, INDArray output) {
+        evaluation.eval(features, output);
+        return 0.0; //Not used
+    }
+
+    @Override
+    protected double scoreMinibatch(Model network, INDArray[] features, INDArray[] labels, INDArray[] fMask, INDArray[] lMask, INDArray[] output) {
+        return scoreMinibatch(network, get0(features), get0(labels), get0(fMask), get0(lMask), get0(output));
+    }
+
+    @Override
+    protected double finalScore(double scoreSum, int minibatchCount, int exampleCount) {
+        return evaluation.scoreForMetric(metric);
+    }
+
+    @Override
+    public boolean minimizeScore() {
+        return metric.minimize();
+    }
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/earlystopping/scorecalc/ClassificationScoreCalculator.java

@@ -0,0 +1,58 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.earlystopping.scorecalc;
+
+import org.deeplearning4j.earlystopping.scorecalc.base.BaseIEvaluationScoreCalculator;
+import org.deeplearning4j.nn.api.Model;
+import org.nd4j.evaluation.classification.Evaluation;
+import org.nd4j.linalg.dataset.api.iterator.DataSetIterator;
+import org.nd4j.linalg.dataset.api.iterator.MultiDataSetIterator;
+
+public class ClassificationScoreCalculator extends BaseIEvaluationScoreCalculator<Model, Evaluation> {
+
+    protected final Evaluation.Metric metric;
+
+    public ClassificationScoreCalculator(Evaluation.Metric metric, DataSetIterator iterator){
+        super(iterator);
+        this.metric = metric;
+    }
+
+    public ClassificationScoreCalculator(Evaluation.Metric metric, MultiDataSetIterator iterator){
+        super(iterator);
+        this.metric = metric;
+    }
+
+    @Override
+    protected Evaluation newEval() {
+        return new Evaluation();
+    }
+
+    @Override
+    protected double finalScore(Evaluation e) {
+        return e.scoreForMetric(metric);
+    }
+
+    @Override
+    public boolean minimizeScore() {
+        //All classification metrics should be maximized: ACCURACY, F1, PRECISION, RECALL, GMEASURE, MCC
+        return false;
+    }
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/earlystopping/scorecalc/DataSetLossCalculator.java

@@ -0,0 +1,115 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.earlystopping.scorecalc;
+
+import org.deeplearning4j.earlystopping.scorecalc.base.BaseScoreCalculator;
+import org.deeplearning4j.nn.api.Model;
+import org.deeplearning4j.nn.graph.ComputationGraph;
+import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;
+import org.nd4j.linalg.api.ndarray.INDArray;
+import org.nd4j.linalg.dataset.DataSet;
+import org.nd4j.linalg.dataset.MultiDataSet;
+import org.nd4j.linalg.dataset.api.iterator.DataSetIterator;
+import org.nd4j.linalg.dataset.api.iterator.MultiDataSetIterator;
+import org.nd4j.shade.jackson.annotation.JsonProperty;
+
+public class DataSetLossCalculator extends BaseScoreCalculator<Model> {
+
+    @JsonProperty
+    private boolean average;
+
+    /**
+     * Calculate the score (loss function value) on a given data set (usually a test set)
+     *
+     * @param dataSetIterator Data set to calculate the score for
+     * @param average         Whether to return the average (sum of loss / N) or just (sum of loss)
+     */
+    public DataSetLossCalculator(DataSetIterator dataSetIterator, boolean average) {
+        super(dataSetIterator);
+        this.average = average;
+    }
+
+    /**Calculate the score (loss function value) on a given data set (usually a test set)
+     *
+     * @param dataSetIterator Data set to calculate the score for
+     * @param average Whether to return the average (sum of loss / N) or just (sum of loss)
+     */
+    public DataSetLossCalculator(MultiDataSetIterator dataSetIterator, boolean average) {
+        super(dataSetIterator);
+        this.average = average;
+    }
+
+    @Override
+    public String toString() {
+        return "DataSetLossCalculator(average=" + average + ")";
+    }
+
+    @Override
+    protected void reset() {
+        scoreSum = 0;
+        minibatchCount = 0;
+        exampleCount = 0;
+    }
+
+    @Override
+    protected INDArray output(Model network, INDArray input, INDArray fMask, INDArray lMask) {
+        return output(network, arr(input), arr(fMask), arr(lMask))[0];
+    }
+
+    @Override
+    protected INDArray[] output(Model network, INDArray[] input, INDArray[] fMask, INDArray[] lMask) {
+        if(network instanceof MultiLayerNetwork){
+            INDArray out = ((MultiLayerNetwork) network).output(input[0], false, get0(fMask), get0(lMask));
+            return new INDArray[]{out};
+        } else if(network instanceof ComputationGraph){
+            return ((ComputationGraph) network).output(false, input, fMask, lMask);
+        } else {
+            throw new RuntimeException("Unknown model type: " + network.getClass());
+        }
+    }
+
+    @Override
+    protected double scoreMinibatch(Model network, INDArray[] features, INDArray[] labels, INDArray[] fMask, INDArray[] lMask, INDArray[] output) {
+        if(network instanceof MultiLayerNetwork){
+            return ((MultiLayerNetwork) network).score(new DataSet(get0(features), get0(labels), get0(fMask), get0(lMask)), false)
+                    * features[0].size(0);
+        } else if(network instanceof ComputationGraph){
+            return ((ComputationGraph) network).score(new MultiDataSet(features, labels, fMask, lMask))
+                    * features[0].size(0);
+        } else {
+            throw new RuntimeException("Unknown model type: " + network.getClass());
+        }
+    }
+
+    @Override
+    protected double finalScore(double scoreSum, int minibatchCount, int exampleCount) {
+        if(average){
+            return scoreSum / exampleCount;
+        } else {
+            return scoreSum;
+        }
+    }
+
+    @Override
+    public boolean minimizeScore() {
+        return true;    //Minimize loss
+    }
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/earlystopping/scorecalc/DataSetLossCalculatorCG.java

@@ -0,0 +1,103 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.earlystopping.scorecalc;
+
+import lombok.NoArgsConstructor;
+import lombok.val;
+import org.deeplearning4j.nn.graph.ComputationGraph;
+import org.nd4j.linalg.dataset.DataSet;
+import org.nd4j.linalg.dataset.api.MultiDataSet;
+import org.nd4j.linalg.dataset.api.iterator.DataSetIterator;
+import org.nd4j.linalg.dataset.api.iterator.MultiDataSetIterator;
+import org.nd4j.shade.jackson.annotation.JsonIgnore;
+import org.nd4j.shade.jackson.annotation.JsonProperty;
+
+@NoArgsConstructor
+@Deprecated
+public class DataSetLossCalculatorCG implements ScoreCalculator<ComputationGraph> {
+    @JsonIgnore
+    private DataSetIterator dataSetIterator;
+    @JsonIgnore
+    private MultiDataSetIterator multiDataSetIterator;
+    @JsonProperty
+    private boolean average;
+
+    /**Calculate the score (loss function value) on a given data set (usually a test set)
+     *
+     * @param dataSetIterator Data set to calculate the score for
+     * @param average Whether to return the average (sum of loss / N) or just (sum of loss)
+     */
+    public DataSetLossCalculatorCG(DataSetIterator dataSetIterator, boolean average) {
+        this.dataSetIterator = dataSetIterator;
+        this.average = average;
+    }
+
+    /**Calculate the score (loss function value) on a given data set (usually a test set)
+     *
+     * @param dataSetIterator Data set to calculate the score for
+     * @param average Whether to return the average (sum of loss / N) or just (sum of loss)
+     */
+    public DataSetLossCalculatorCG(MultiDataSetIterator dataSetIterator, boolean average) {
+        this.multiDataSetIterator = dataSetIterator;
+        this.average = average;
+    }
+
+    @Override
+    public double calculateScore(ComputationGraph network) {
+        double lossSum = 0.0;
+        int exCount = 0;
+
+        if (dataSetIterator != null) {
+            dataSetIterator.reset();
+
+            while (dataSetIterator.hasNext()) {
+                DataSet dataSet = dataSetIterator.next();
+                val nEx = dataSet.getFeatures().size(0);
+                lossSum += network.score(dataSet) * nEx;
+                exCount += nEx;
+            }
+        } else {
+            multiDataSetIterator.reset();
+
+            while (multiDataSetIterator.hasNext()) {
+                MultiDataSet dataSet = multiDataSetIterator.next();
+                val nEx = dataSet.getFeatures(0).size(0);
+                lossSum += network.score(dataSet) * nEx;
+                exCount += nEx;
+            }
+        }
+
+        if (average)
+            return lossSum / exCount;
+        else
+            return lossSum;
+    }
+
+    @Override
+    public boolean minimizeScore() {
+        return true;
+    }
+
+    @Override
+    public String toString() {
+        return "DataSetLossCalculatorCG(" + dataSetIterator + ",average=" + average + ")";
+    }
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/earlystopping/scorecalc/ROCScoreCalculator.java

@@ -0,0 +1,96 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.earlystopping.scorecalc;
+
+import org.deeplearning4j.earlystopping.scorecalc.base.BaseIEvaluationScoreCalculator;
+import org.deeplearning4j.nn.api.Model;
+import org.nd4j.evaluation.IEvaluation;
+import org.nd4j.evaluation.classification.ROC;
+import org.nd4j.evaluation.classification.ROCBinary;
+import org.nd4j.evaluation.classification.ROCMultiClass;
+import org.nd4j.linalg.dataset.api.iterator.DataSetIterator;
+import org.nd4j.linalg.dataset.api.iterator.MultiDataSetIterator;
+
+public class ROCScoreCalculator extends BaseIEvaluationScoreCalculator<Model, IEvaluation> {
+
+    public enum ROCType {ROC, BINARY, MULTICLASS}
+    public enum Metric {AUC, AUPRC};
+
+    protected final ROCType type;
+    protected final Metric metric;
+
+    public ROCScoreCalculator(ROCType type, DataSetIterator iterator) {
+        this(type, Metric.AUC, iterator);
+    }
+
+    public ROCScoreCalculator(ROCType type, MultiDataSetIterator iterator){
+        this(type, Metric.AUC, iterator);
+    }
+
+    public ROCScoreCalculator(ROCType type, Metric metric, DataSetIterator iterator){
+        super(iterator);
+        this.type = type;
+        this.metric = metric;
+    }
+
+    public ROCScoreCalculator(ROCType type, Metric metric, MultiDataSetIterator iterator){
+        super(iterator);
+        this.type = type;
+        this.metric = metric;
+    }
+
+
+    @Override
+    protected IEvaluation newEval() {
+        switch (type){
+            case ROC:
+                return new ROC();
+            case BINARY:
+                return new ROCBinary();
+            case MULTICLASS:
+                return new ROCMultiClass();
+            default:
+                throw new IllegalStateException("Unknown type: " + type);
+        }
+    }
+
+    @Override
+    protected double finalScore(IEvaluation eval) {
+        switch (type){
+            case ROC:
+                ROC r = (ROC)eval;
+                return metric == Metric.AUC ? r.calculateAUC() : r.calculateAUCPR();
+            case BINARY:
+                ROCBinary r2 = (ROCBinary) eval;
+                return metric == Metric.AUC ? r2.calculateAverageAuc() : r2.calculateAverageAUCPR();
+            case MULTICLASS:
+                ROCMultiClass r3 = (ROCMultiClass)eval;
+                return metric == Metric.AUC ? r3.calculateAverageAUC() : r3.calculateAverageAUCPR();
+            default:
+                throw new IllegalStateException("Unknown type: " + type);
+        }
+    }
+
+    @Override
+    public boolean minimizeScore() {
+        return false;   //Maximize AUC, AUPRC
+    }
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/earlystopping/scorecalc/RegressionScoreCalculator.java

@@ -0,0 +1,52 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.earlystopping.scorecalc;
+
+import org.deeplearning4j.earlystopping.scorecalc.base.BaseIEvaluationScoreCalculator;
+import org.deeplearning4j.nn.api.Model;
+import org.nd4j.evaluation.regression.RegressionEvaluation;
+import org.nd4j.evaluation.regression.RegressionEvaluation.Metric;
+import org.nd4j.linalg.dataset.api.iterator.DataSetIterator;
+
+public class RegressionScoreCalculator extends BaseIEvaluationScoreCalculator<Model, RegressionEvaluation> {
+
+    protected final Metric metric;
+
+    public RegressionScoreCalculator(Metric metric, DataSetIterator iterator){
+        super(iterator);
+        this.metric = metric;
+    }
+
+    @Override
+    protected RegressionEvaluation newEval() {
+        return new RegressionEvaluation();
+    }
+
+    @Override
+    protected double finalScore(RegressionEvaluation eval) {
+        return eval.scoreForMetric(metric);
+    }
+
+    @Override
+    public boolean minimizeScore() {
+        return metric.minimize();
+    }
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/earlystopping/scorecalc/ScoreCalculator.java

@@ -0,0 +1,46 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.earlystopping.scorecalc;
+
+import org.deeplearning4j.nn.api.Model;
+import org.nd4j.shade.jackson.annotation.JsonInclude;
+import org.nd4j.shade.jackson.annotation.JsonSubTypes;
+import org.nd4j.shade.jackson.annotation.JsonTypeInfo;
+
+import java.io.Serializable;
+
+@JsonTypeInfo(use = JsonTypeInfo.Id.CLASS, include = JsonTypeInfo.As.PROPERTY, property = "@class")
+@JsonInclude(JsonInclude.Include.NON_NULL)
+@JsonSubTypes(value = {
+                @JsonSubTypes.Type(value = DataSetLossCalculator.class, name = "BestScoreEpochTerminationCondition"),
+                @JsonSubTypes.Type(value = DataSetLossCalculatorCG.class, name = "MaxEpochsTerminationCondition"),
+
+})
+public interface ScoreCalculator<T extends Model> extends Serializable {
+
+    /** Calculate the score for the given MultiLayerNetwork */
+    double calculateScore(T network);
+
+    /**
+     * @return If true: the score should be minimized. If false: the score should be maximized.
+     */
+    boolean minimizeScore();
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/earlystopping/scorecalc/VAEReconErrorScoreCalculator.java

@@ -0,0 +1,102 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.earlystopping.scorecalc;
+
+import org.deeplearning4j.earlystopping.scorecalc.base.BaseScoreCalculator;
+import org.deeplearning4j.nn.api.Layer;
+import org.deeplearning4j.nn.api.Model;
+import org.deeplearning4j.nn.graph.ComputationGraph;
+import org.deeplearning4j.nn.layers.variational.VariationalAutoencoder;
+import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;
+import org.nd4j.evaluation.regression.RegressionEvaluation;
+import org.nd4j.evaluation.regression.RegressionEvaluation.Metric;
+import org.nd4j.linalg.api.ndarray.INDArray;
+import org.nd4j.linalg.dataset.api.iterator.DataSetIterator;
+import org.deeplearning4j.nn.workspace.LayerWorkspaceMgr;
+
+public class VAEReconErrorScoreCalculator extends BaseScoreCalculator<Model> {
+
+    protected final Metric metric;
+    protected RegressionEvaluation evaluation;
+
+    /**
+     * Constructor for reconstruction *ERROR*
+     *
+     * @param metric
+     * @param iterator
+     */
+    public VAEReconErrorScoreCalculator(Metric metric, DataSetIterator iterator) {
+        super(iterator);
+        this.metric = metric;
+    }
+
+    @Override
+    protected void reset() {
+        evaluation = new RegressionEvaluation();
+    }
+
+    @Override
+    protected INDArray output(Model net, INDArray input, INDArray fMask, INDArray lMask) {
+        Layer l;
+        if(net instanceof MultiLayerNetwork) {
+            MultiLayerNetwork network = (MultiLayerNetwork)net;
+            l = network.getLayer(0);
+        } else {
+            ComputationGraph network = (ComputationGraph)net;
+            l = network.getLayer(0);
+        }
+
+        if(!(l instanceof VariationalAutoencoder)){
+            throw new UnsupportedOperationException("Can only score networks with VariationalAutoencoder layers as first layer -" +
+                    " got " + l.getClass().getSimpleName());
+        }
+        VariationalAutoencoder vae = (VariationalAutoencoder)l;
+        INDArray z = vae.activate(input, false, LayerWorkspaceMgr.noWorkspaces());
+        return vae.generateAtMeanGivenZ(z);
+    }
+
+    @Override
+    protected INDArray[] output(Model network, INDArray[] input, INDArray[] fMask, INDArray[] lMask) {
+        return new INDArray[]{output(network, get0(input), get0(fMask), get0(lMask))};
+    }
+
+    @Override
+    protected double scoreMinibatch(Model network, INDArray features, INDArray labels, INDArray fMask,
+                                    INDArray lMask, INDArray output) {
+        evaluation.eval(features, output);
+        return 0.0; //Not used
+    }
+
+    @Override
+    protected double scoreMinibatch(Model network, INDArray[] features, INDArray[] labels, INDArray[] fMask, INDArray[] lMask, INDArray[] output) {
+        return scoreMinibatch(network, get0(features), get0(labels), get0(fMask), get0(lMask), get0(output));
+    }
+
+    @Override
+    protected double finalScore(double scoreSum, int minibatchCount, int exampleCount) {
+        return evaluation.scoreForMetric(metric);
+    }
+
+    @Override
+    public boolean minimizeScore() {
+        return true;    //Minimize reconstruction error
+    }
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/earlystopping/scorecalc/VAEReconProbScoreCalculator.java

@@ -0,0 +1,128 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.earlystopping.scorecalc;
+
+import org.deeplearning4j.earlystopping.scorecalc.base.BaseScoreCalculator;
+import org.deeplearning4j.nn.api.Layer;
+import org.deeplearning4j.nn.api.Model;
+import org.deeplearning4j.nn.graph.ComputationGraph;
+import org.deeplearning4j.nn.layers.variational.VariationalAutoencoder;
+import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;
+import org.nd4j.linalg.api.ndarray.INDArray;
+import org.nd4j.linalg.dataset.api.iterator.DataSetIterator;
+
+public class VAEReconProbScoreCalculator extends BaseScoreCalculator<Model> {
+
+    protected final int reconstructionProbNumSamples;
+    protected final boolean logProb;
+    protected final boolean average;
+
+    /**
+     * Constructor for average reconstruction probability
+     *
+     * @param iterator Iterator
+     * @param reconstructionProbNumSamples Number of samples. See {@link VariationalAutoencoder#reconstructionProbability(INDArray, int)}
+     *                                    for details
+     * @param logProb If true: calculate (negative) log probability. False: probability
+     */
+    public VAEReconProbScoreCalculator(DataSetIterator iterator, int reconstructionProbNumSamples, boolean logProb) {
+        this(iterator, reconstructionProbNumSamples, logProb, true);
+    }
+
+    /**
+     * Constructor for reconstruction probability
+     *
+     * @param iterator Iterator
+     * @param reconstructionProbNumSamples Number of samples. See {@link VariationalAutoencoder#reconstructionProbability(INDArray, int)}
+     *                                    for details
+     * @param logProb If true: calculate (negative) log probability. False: probability
+     * @param average If true: return average (log) probability. False: sum of log probability.
+     *
+     */
+    public VAEReconProbScoreCalculator(DataSetIterator iterator, int reconstructionProbNumSamples, boolean logProb,
+                                       boolean average){
+        super(iterator);
+        this.reconstructionProbNumSamples = reconstructionProbNumSamples;
+        this.logProb = logProb;
+        this.average = average;
+    }
+
+    @Override
+    protected void reset() {
+        scoreSum = 0;
+        minibatchCount = 0;
+        exampleCount = 0;
+    }
+
+    @Override
+    protected INDArray output(Model network, INDArray input, INDArray fMask, INDArray lMask) {
+        return null;    //Not used
+    }
+
+    @Override
+    protected INDArray[] output(Model network, INDArray[] input, INDArray[] fMask, INDArray[] lMask) {
+        return null;    //Not used
+    }
+
+    @Override
+    protected double scoreMinibatch(Model net, INDArray features, INDArray labels, INDArray fMask,
+                                    INDArray lMask, INDArray output) {
+        Layer l;
+        if(net instanceof MultiLayerNetwork) {
+            MultiLayerNetwork network = (MultiLayerNetwork)net;
+            l = network.getLayer(0);
+        } else {
+            ComputationGraph network = (ComputationGraph)net;
+            l = network.getLayer(0);
+        }
+
+        if(!(l instanceof VariationalAutoencoder)) {
+            throw new UnsupportedOperationException("Can only score networks with VariationalAutoencoder layers as first layer -" +
+                    " got " + l.getClass().getSimpleName());
+        }
+        VariationalAutoencoder vae = (VariationalAutoencoder)l;
+        //Reconstruction prob
+        if(logProb) {
+            return -vae.reconstructionLogProbability(features, reconstructionProbNumSamples).sumNumber().doubleValue();
+        } else {
+            return vae.reconstructionProbability(features, reconstructionProbNumSamples).sumNumber().doubleValue();
+        }
+    }
+
+    @Override
+    protected double scoreMinibatch(Model network, INDArray[] features, INDArray[] labels, INDArray[] fMask, INDArray[] lMask, INDArray[] output) {
+        return 0;
+    }
+
+    @Override
+    protected double finalScore(double scoreSum, int minibatchCount, int exampleCount) {
+        if(average){
+            return scoreSum / exampleCount;
+        } else {
+            return scoreSum;
+        }
+    }
+
+    @Override
+    public boolean minimizeScore() {
+        return false;   //Maximize the reconstruction probability
+    }
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/earlystopping/scorecalc/base/BaseIEvaluationScoreCalculator.java

@@ -0,0 +1,67 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.earlystopping.scorecalc.base;
+
+import org.deeplearning4j.datasets.iterator.MultiDataSetWrapperIterator;
+import org.deeplearning4j.earlystopping.scorecalc.ScoreCalculator;
+import org.deeplearning4j.nn.api.Model;
+import org.deeplearning4j.nn.graph.ComputationGraph;
+import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;
+import org.nd4j.evaluation.IEvaluation;
+import org.nd4j.linalg.dataset.adapter.MultiDataSetIteratorAdapter;
+import org.nd4j.linalg.dataset.api.iterator.DataSetIterator;
+import org.nd4j.linalg.dataset.api.iterator.MultiDataSetIterator;
+
+public abstract class BaseIEvaluationScoreCalculator<T extends Model, U extends IEvaluation> implements ScoreCalculator<T> {
+
+    protected MultiDataSetIterator iterator;
+    protected DataSetIterator iter;
+
+    protected BaseIEvaluationScoreCalculator(MultiDataSetIterator iterator){
+        this.iterator = iterator;
+    }
+
+    protected BaseIEvaluationScoreCalculator(DataSetIterator iterator){
+        this.iter = iterator;
+    }
+
+    @Override
+    public double calculateScore(T network) {
+        U eval = newEval();
+
+        if(network instanceof MultiLayerNetwork){
+            DataSetIterator i = (iter != null ? iter : new MultiDataSetWrapperIterator(iterator));
+            eval = ((MultiLayerNetwork) network).doEvaluation(i, eval)[0];
+        } else if(network instanceof ComputationGraph){
+            MultiDataSetIterator i = (iterator != null ? iterator : new MultiDataSetIteratorAdapter(iter));
+            eval = ((ComputationGraph) network).doEvaluation(i, eval)[0];
+        } else {
+            throw new RuntimeException("Unknown model type: " + network.getClass());
+        }
+        return finalScore(eval);
+    }
+
+    protected abstract U newEval();
+
+    protected abstract double finalScore(U eval);
+
+
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/earlystopping/scorecalc/base/BaseMLNScoreCalculator.java

@@ -0,0 +1,49 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.earlystopping.scorecalc.base;
+
+import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;
+import org.nd4j.linalg.api.ndarray.INDArray;
+import org.nd4j.linalg.dataset.api.iterator.DataSetIterator;
+
+public abstract class BaseMLNScoreCalculator extends BaseScoreCalculator<MultiLayerNetwork> {
+
+
+    protected BaseMLNScoreCalculator(DataSetIterator iterator) {
+        super(iterator);
+    }
+
+    @Override
+    protected INDArray output(MultiLayerNetwork network, INDArray input, INDArray fMask, INDArray lMask) {
+        return network.output(input, false, fMask, lMask);
+    }
+
+    @Override
+    protected double scoreMinibatch(MultiLayerNetwork network, INDArray[] features, INDArray[] labels, INDArray[] fMask,
+                                    INDArray[] lMask, INDArray[] output) {
+        return scoreMinibatch(network, get0(features), get0(labels), get0(fMask), get0(lMask), get0(output));
+    }
+
+    @Override
+    protected INDArray[] output(MultiLayerNetwork network, INDArray[] input, INDArray[] fMask, INDArray[] lMask) {
+        return new INDArray[]{output(network, get0(input), get0(fMask), get0(lMask))};
+    }
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/earlystopping/scorecalc/base/BaseScoreCalculator.java

@@ -0,0 +1,109 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.earlystopping.scorecalc.base;
+
+import lombok.NonNull;
+import org.deeplearning4j.earlystopping.scorecalc.ScoreCalculator;
+import org.deeplearning4j.nn.api.Model;
+import org.nd4j.linalg.api.ndarray.INDArray;
+import org.nd4j.linalg.dataset.DataSet;
+import org.nd4j.linalg.dataset.api.MultiDataSet;
+import org.nd4j.linalg.dataset.api.iterator.DataSetIterator;
+import org.nd4j.linalg.dataset.api.iterator.MultiDataSetIterator;
+
+public abstract class BaseScoreCalculator<T extends Model> implements ScoreCalculator<T> {
+
+    protected MultiDataSetIterator mdsIterator;
+    protected DataSetIterator iterator;
+    protected double scoreSum;
+    protected int minibatchCount;
+    protected int exampleCount;
+
+    protected BaseScoreCalculator(@NonNull DataSetIterator iterator){
+        this.iterator = iterator;
+    }
+
+    protected BaseScoreCalculator(@NonNull MultiDataSetIterator iterator){
+        this.mdsIterator = iterator;
+    }
+
+    @Override
+    public double calculateScore(T network) {
+        reset();
+
+        if(iterator != null) {
+            if (!iterator.hasNext())
+                iterator.reset();
+
+            while (iterator.hasNext()) {
+                DataSet ds = iterator.next();
+                INDArray out = output(network, ds.getFeatures(), ds.getFeaturesMaskArray(), ds.getLabelsMaskArray());
+                scoreSum += scoreMinibatch(network, ds.getFeatures(), ds.getLabels(), ds.getFeaturesMaskArray(),
+                        ds.getLabelsMaskArray(), out);
+                minibatchCount++;
+                exampleCount += ds.getFeatures().size(0);
+            }
+        } else {
+            if(!mdsIterator.hasNext())
+                mdsIterator.reset();
+
+            while(mdsIterator.hasNext()){
+                MultiDataSet mds = mdsIterator.next();
+                INDArray[] out = output(network, mds.getFeatures(), mds.getFeaturesMaskArrays(), mds.getLabelsMaskArrays() );
+                scoreSum += scoreMinibatch(network, mds.getFeatures(), mds.getLabels(), mds.getFeaturesMaskArrays(),
+                        mds.getLabelsMaskArrays(), out);
+                minibatchCount++;
+                exampleCount += mds.getFeatures(0).size(0);
+            }
+        }
+
+        return finalScore(scoreSum, minibatchCount, exampleCount);
+    }
+
+    protected abstract void reset();
+
+    protected abstract INDArray output(T network, INDArray input, INDArray fMask, INDArray lMask);
+
+    protected abstract INDArray[] output(T network, INDArray[] input, INDArray[] fMask, INDArray[] lMask);
+
+    protected double scoreMinibatch(T network, INDArray features, INDArray labels,
+                                    INDArray fMask, INDArray lMask, INDArray output){
+        return scoreMinibatch(network, arr(features), arr(labels), arr(fMask), arr(lMask), arr(output));
+    }
+
+    protected abstract double scoreMinibatch(T network, INDArray[] features, INDArray[] labels,
+                                             INDArray[] fMask, INDArray[] lMask, INDArray[] output);
+
+    protected abstract double finalScore(double scoreSum, int minibatchCount, int exampleCount);
+
+    public static INDArray[] arr(INDArray in){
+        if(in == null) return null;
+        return new INDArray[]{in};
+    }
+
+    public static INDArray get0(INDArray[] in){
+        if(in == null) return null;
+        if(in.length != 1){
+            throw new IllegalStateException("Expected length 1 array here: got length " + in.length);
+        }
+        return in[0];
+    }
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/earlystopping/termination/BestScoreEpochTerminationCondition.java

@@ -0,0 +1,61 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.earlystopping.termination;
+
+import lombok.Data;
+import org.nd4j.shade.jackson.annotation.JsonProperty;
+
+@Data
+public class BestScoreEpochTerminationCondition implements EpochTerminationCondition {
+    @JsonProperty
+    private final double bestExpectedScore;
+
+    public BestScoreEpochTerminationCondition(@JsonProperty("bestExpectedScore") double bestExpectedScore) {
+        this.bestExpectedScore = bestExpectedScore;
+    }
+
+    /**
+     * @deprecated "lessBetter" argument no longer used
+     */
+    @Deprecated
+    public BestScoreEpochTerminationCondition(double bestExpectedScore, boolean lesserBetter) {
+        this(bestExpectedScore);
+    }
+
+    @Override
+    public void initialize() {
+        /* No OP */
+    }
+
+    @Override
+    public boolean terminate(int epochNum, double score, boolean minimize) {
+        if (minimize) {
+            return score < bestExpectedScore;
+        } else {
+            return bestExpectedScore < score;
+        }
+    }
+
+    @Override
+    public String toString() {
+        return "BestScoreEpochTerminationCondition(" + bestExpectedScore + ")";
+    }
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/earlystopping/termination/EpochTerminationCondition.java

@@ -0,0 +1,45 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.earlystopping.termination;
+
+
+import org.nd4j.shade.jackson.annotation.JsonInclude;
+import org.nd4j.shade.jackson.annotation.JsonSubTypes;
+import org.nd4j.shade.jackson.annotation.JsonTypeInfo;
+
+import java.io.Serializable;
+
+@JsonTypeInfo(use = JsonTypeInfo.Id.CLASS, include = JsonTypeInfo.As.PROPERTY, property = "@class")
+@JsonInclude(JsonInclude.Include.NON_NULL)
+public interface EpochTerminationCondition extends Serializable {
+
+    /** Initialize the epoch termination condition (often a no-op)*/
+    void initialize();
+
+    /**Should the early stopping training terminate at this epoch, based on the calculated score and the epoch number?
+     * Returns true if training should terminated, or false otherwise
+     * @param epochNum Number of the last completed epoch (starting at 0)
+     * @param score Score calculate for this epoch
+     * @return Whether training should be terminated at this epoch
+     */
+    boolean terminate(int epochNum, double score, boolean minimize);
+
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/earlystopping/termination/InvalidScoreIterationTerminationCondition.java

@@ -0,0 +1,41 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.earlystopping.termination;
+
+import lombok.Data;
+
+@Data
+public class InvalidScoreIterationTerminationCondition implements IterationTerminationCondition {
+    @Override
+    public void initialize() {
+        //No op
+    }
+
+    @Override
+    public boolean terminate(double lastMiniBatchScore) {
+        return Double.isNaN(lastMiniBatchScore) || Double.isInfinite(lastMiniBatchScore);
+    }
+
+    @Override
+    public String toString() {
+        return "InvalidScoreIterationTerminationCondition()";
+    }
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/earlystopping/termination/IterationTerminationCondition.java

@@ -0,0 +1,42 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.earlystopping.termination;
+
+import org.nd4j.shade.jackson.annotation.JsonInclude;
+import org.nd4j.shade.jackson.annotation.JsonTypeInfo;
+
+import java.io.Serializable;
+
+@JsonTypeInfo(use = JsonTypeInfo.Id.CLASS, include = JsonTypeInfo.As.PROPERTY, property = "@class")
+@JsonInclude(JsonInclude.Include.NON_NULL)
+public interface IterationTerminationCondition extends Serializable {
+
+    /** Initialize the iteration termination condition (sometimes a no-op)*/
+    void initialize();
+
+    /** Should early stopping training terminate at this iteration, based on the score for the last iteration?
+     * return true if training should be terminated immediately, or false otherwise
+     * @param lastMiniBatchScore Score of the last minibatch
+     * @return whether to terminate or not
+     */
+    boolean terminate(double lastMiniBatchScore);
+
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/earlystopping/termination/MaxEpochsTerminationCondition.java

@@ -0,0 +1,55 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.earlystopping.termination;
+
+import lombok.Data;
+import lombok.NoArgsConstructor;
+import org.nd4j.shade.jackson.annotation.JsonCreator;
+import org.nd4j.shade.jackson.annotation.JsonProperty;
+
+@NoArgsConstructor
+@Data
+public class MaxEpochsTerminationCondition implements EpochTerminationCondition {
+    @JsonProperty
+    private int maxEpochs;
+
+    @JsonCreator
+    public MaxEpochsTerminationCondition(int maxEpochs) {
+        if (maxEpochs <= 0)
+            throw new IllegalArgumentException("Max number of epochs must be >= 1");
+        this.maxEpochs = maxEpochs;
+    }
+
+    @Override
+    public void initialize() {
+        //No op
+    }
+
+    @Override
+    public boolean terminate(int epochNum, double score, boolean minimize) {
+        return epochNum + 1 >= maxEpochs; //epochNum starts at 0
+    }
+
+    @Override
+    public String toString() {
+        return "MaxEpochsTerminationCondition(" + maxEpochs + ")";
+    }
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/earlystopping/termination/MaxScoreIterationTerminationCondition.java

@@ -0,0 +1,49 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.earlystopping.termination;
+
+import lombok.Data;
+import org.nd4j.shade.jackson.annotation.JsonProperty;
+
+@Data
+public class MaxScoreIterationTerminationCondition implements IterationTerminationCondition {
+
+    private double maxScore;
+
+    public MaxScoreIterationTerminationCondition(@JsonProperty("maxScore") double maxScore) {
+        this.maxScore = maxScore;
+    }
+
+    @Override
+    public void initialize() {
+        //no op
+    }
+
+    @Override
+    public boolean terminate(double lastMiniBatchScore) {
+        return lastMiniBatchScore > maxScore || Double.isNaN(lastMiniBatchScore);
+    }
+
+    @Override
+    public String toString() {
+        return "MaxScoreIterationTerminationCondition(" + maxScore + ")";
+    }
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/earlystopping/termination/MaxTimeIterationTerminationCondition.java

@@ -0,0 +1,61 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.earlystopping.termination;
+
+import lombok.Data;
+import org.nd4j.shade.jackson.annotation.JsonProperty;
+
+import java.util.concurrent.TimeUnit;
+
+/**Terminate training based on max time.
+ */
+@Data
+public class MaxTimeIterationTerminationCondition implements IterationTerminationCondition {
+
+    private long maxTimeAmount;
+    private TimeUnit maxTimeUnit;
+    private long initializationTime;
+    private long endTime;
+
+    public MaxTimeIterationTerminationCondition(@JsonProperty("maxTimeAmount") long maxTimeAmount, @JsonProperty("maxTimeUnit") TimeUnit maxTimeUnit) {
+        if (maxTimeAmount <= 0 || maxTimeUnit == null)
+            throw new IllegalArgumentException(
+                            "Invalid maximum training time: " + "amount = " + maxTimeAmount + " unit = " + maxTimeUnit);
+        this.maxTimeAmount = maxTimeAmount;
+        this.maxTimeUnit = maxTimeUnit;
+    }
+
+    @Override
+    public void initialize() {
+        initializationTime = System.currentTimeMillis();
+        endTime = initializationTime + maxTimeUnit.toMillis(maxTimeAmount);
+    }
+
+    @Override
+    public boolean terminate(double lastMiniBatchScore) {
+        return System.currentTimeMillis() >= endTime;
+    }
+
+    @Override
+    public String toString() {
+        return "MaxTimeIterationTerminationCondition(" + maxTimeAmount + ",unit=" + maxTimeUnit + ")";
+    }
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/earlystopping/termination/ScoreImprovementEpochTerminationCondition.java

@@ -0,0 +1,81 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.earlystopping.termination;
+
+import lombok.Data;
+import lombok.extern.slf4j.Slf4j;
+import org.nd4j.shade.jackson.annotation.JsonProperty;
+
+@Slf4j
+@Data
+public class ScoreImprovementEpochTerminationCondition implements EpochTerminationCondition {
+    @JsonProperty
+    private int maxEpochsWithNoImprovement;
+    @JsonProperty
+    private int bestEpoch = -1;
+    @JsonProperty
+    private double bestScore;
+    @JsonProperty
+    private double minImprovement = 0.0;
+
+    public ScoreImprovementEpochTerminationCondition(int maxEpochsWithNoImprovement) {
+        this.maxEpochsWithNoImprovement = maxEpochsWithNoImprovement;
+    }
+
+    public ScoreImprovementEpochTerminationCondition(@JsonProperty("maxEpochsWithNoImprovement") int maxEpochsWithNoImprovement,
+                                                     @JsonProperty("minImprovement") double minImprovement) {
+        this.maxEpochsWithNoImprovement = maxEpochsWithNoImprovement;
+        this.minImprovement = minImprovement;
+    }
+
+    @Override
+    public void initialize() {
+        bestEpoch = -1;
+        bestScore = Double.NaN;
+    }
+
+    @Override
+    public boolean terminate(int epochNum, double score, boolean minimize) {
+        if (bestEpoch == -1) {
+            bestEpoch = epochNum;
+            bestScore = score;
+            return false;
+        } else {
+            double improvement = (minimize ? bestScore - score : score - bestScore);
+            if (improvement > minImprovement) {
+                if (minImprovement > 0) {
+                    log.info("Epoch with score greater than threshold * * *");
+                }
+                bestScore = score;
+                bestEpoch = epochNum;
+                return false;
+            }
+
+            return epochNum >= bestEpoch + maxEpochsWithNoImprovement;
+        }
+    }
+
+    @Override
+    public String toString() {
+        return "ScoreImprovementEpochTerminationCondition(maxEpochsWithNoImprovement=" + maxEpochsWithNoImprovement
+                        + ", minImprovement=" + minImprovement + ")";
+    }
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/earlystopping/trainer/BaseEarlyStoppingTrainer.java

@@ -0,0 +1,385 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.earlystopping.trainer;
+
+import org.deeplearning4j.earlystopping.EarlyStoppingConfiguration;
+import org.deeplearning4j.earlystopping.EarlyStoppingResult;
+import org.deeplearning4j.earlystopping.listener.EarlyStoppingListener;
+import org.deeplearning4j.earlystopping.scorecalc.ScoreCalculator;
+import org.deeplearning4j.earlystopping.termination.EpochTerminationCondition;
+import org.deeplearning4j.earlystopping.termination.IterationTerminationCondition;
+import org.deeplearning4j.nn.api.Model;
+import org.deeplearning4j.nn.graph.ComputationGraph;
+import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;
+import org.deeplearning4j.optimize.api.TrainingListener;
+import org.nd4j.common.base.Preconditions;
+import org.nd4j.linalg.dataset.AsyncDataSetIterator;
+import org.nd4j.linalg.dataset.AsyncMultiDataSetIterator;
+import org.nd4j.linalg.dataset.DataSet;
+import org.nd4j.linalg.dataset.api.MultiDataSet;
+import org.nd4j.linalg.dataset.api.iterator.DataSetIterator;
+import org.nd4j.linalg.dataset.api.iterator.MultiDataSetIterator;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.FileNotFoundException;
+import java.io.IOException;
+import java.util.Collection;
+import java.util.Iterator;
+import java.util.LinkedHashMap;
+import java.util.Map;
+
+;
+
+public abstract class BaseEarlyStoppingTrainer<T extends Model> implements IEarlyStoppingTrainer<T> {
+
+    private static Logger log = LoggerFactory.getLogger(BaseEarlyStoppingTrainer.class);
+
+    protected T model;
+
+    protected final EarlyStoppingConfiguration<T> esConfig;
+    private final DataSetIterator train;
+    private final MultiDataSetIterator trainMulti;
+    private final Iterator<?> iterator;
+    private EarlyStoppingListener<T> listener;
+
+    private double bestModelScore = Double.MAX_VALUE;
+    private int bestModelEpoch = -1;
+
+    protected BaseEarlyStoppingTrainer(EarlyStoppingConfiguration<T> earlyStoppingConfiguration, T model,
+                                       DataSetIterator train, MultiDataSetIterator trainMulti, EarlyStoppingListener<T> listener) {
+        if(train != null && train.asyncSupported()){
+            train = new AsyncDataSetIterator(train);
+        }
+        if(trainMulti != null && trainMulti.asyncSupported()){
+            trainMulti = new AsyncMultiDataSetIterator(trainMulti);
+        }
+
+        this.esConfig = earlyStoppingConfiguration;
+        this.model = model;
+        this.train = train;
+        this.trainMulti = trainMulti;
+        this.iterator = (train != null ? train : trainMulti);
+        this.listener = listener;
+    }
+
+    protected abstract void fit(DataSet ds);
+
+    protected abstract void fit(MultiDataSet mds);
+
+    protected abstract void pretrain(DataSet ds);
+
+    protected abstract void pretrain(MultiDataSet mds);
+
+    @Override
+    public EarlyStoppingResult<T> fit() {
+        return fit(false);
+    }
+
+    @Override
+    public EarlyStoppingResult<T> pretrain(){
+        return fit(true);
+    }
+
+    protected EarlyStoppingResult<T> fit(boolean pretrain) {
+        esConfig.validate();
+        log.info("Starting early stopping training");
+        if (esConfig.getScoreCalculator() == null)
+            log.warn("No score calculator provided for early stopping. Score will be reported as 0.0 to epoch termination conditions");
+
+        //Initialize termination conditions:
+        if (esConfig.getIterationTerminationConditions() != null) {
+            for (IterationTerminationCondition c : esConfig.getIterationTerminationConditions()) {
+                c.initialize();
+            }
+        }
+        if (esConfig.getEpochTerminationConditions() != null) {
+            for (EpochTerminationCondition c : esConfig.getEpochTerminationConditions()) {
+                c.initialize();
+            }
+        }
+
+        if (listener != null) {
+            listener.onStart(esConfig, model);
+        }
+
+        Map<Integer, Double> scoreVsEpoch = new LinkedHashMap<>();
+
+        Preconditions.checkNotNull(esConfig.getScoreCalculator(), "Score calculator cannot be null");
+        if(esConfig.getScoreCalculator().minimizeScore()){
+            bestModelScore = Double.MAX_VALUE;
+        } else {
+            bestModelScore = -Double.MAX_VALUE;
+        }
+
+        int epochCount = 0;
+        while (true) {
+            reset();
+            double lastScore;
+            boolean terminate = false;
+            IterationTerminationCondition terminationReason = null;
+            int iterCount = 0;
+            triggerEpochListeners(true, model, epochCount);
+            while (iterator.hasNext()) {
+                try {
+                    if(pretrain) {
+                        if(train != null) {
+                            pretrain((DataSet)iterator.next());
+                        } else {
+                            pretrain(trainMulti.next());
+                        }
+                    } else {
+                        if (train != null) {
+                            fit((DataSet) iterator.next());
+                        } else
+                            fit(trainMulti.next());
+                    }
+                } catch (Exception e) {
+                    log.warn("Early stopping training terminated due to exception at epoch {}, iteration {}",
+                            epochCount, iterCount, e);
+                    //Load best model to return
+                    T bestModel;
+                    try {
+                        bestModel = esConfig.getModelSaver().getBestModel();
+
+                       if(bestModel != null)
+                        bestModelScore = bestModel.score();
+                        
+                    } catch (IOException e2) {
+                        throw new RuntimeException(e2);
+                    }
+                    return new EarlyStoppingResult<>(EarlyStoppingResult.TerminationReason.Error, e.toString(),
+                            scoreVsEpoch, bestModelEpoch, bestModelScore, epochCount, bestModel);
+                }
+
+                //Check per-iteration termination conditions
+                if(pretrain){
+                    //TODO support for non-first-layer pretraining
+                    if(model instanceof MultiLayerNetwork) {
+                        lastScore = (((MultiLayerNetwork) model).getLayer(0)).score();
+                        ((MultiLayerNetwork) model).setScore(lastScore);
+                    } else {
+                        ComputationGraph computationGraph = (ComputationGraph) model;
+                        lastScore = computationGraph.getLayer(0).score();
+                        computationGraph.setScore(lastScore);
+                    }
+                } else {
+                    lastScore = model.score();
+                }
+                for (IterationTerminationCondition c : esConfig.getIterationTerminationConditions()) {
+                    if (c.terminate(lastScore)) {
+                        terminate = true;
+                        terminationReason = c;
+                        break;
+                    }
+                }
+                if (terminate) {
+                    break;
+                }
+
+                iterCount++;
+            }
+
+            if(!iterator.hasNext()){
+                //End of epoch (if iterator does have next - means terminated)
+                triggerEpochListeners(false, model, epochCount);
+            }
+
+            if (terminate) {
+                //Handle termination condition:
+                log.info("Hit per iteration epoch termination condition at epoch {}, iteration {}. Reason: {}",
+                        epochCount, iterCount, terminationReason);
+
+                if (esConfig.isSaveLastModel()) {
+                    //Save last model:
+                    try {
+                        esConfig.getModelSaver().saveLatestModel(model, 0.0);
+                    } catch (IOException e) {
+                        //best model not saved, let's just use default
+                        if(e instanceof FileNotFoundException) {
+
+                        }
+                        else
+                            throw new RuntimeException("Error saving most recent model", e);
+                    }
+                }
+
+                T bestModel;
+                try {
+                    bestModel = esConfig.getModelSaver().getBestModel();
+                } catch (IOException e2) {
+                    throw new RuntimeException(e2);
+                }
+
+
+                EarlyStoppingResult<T> result = new EarlyStoppingResult<>(
+                        EarlyStoppingResult.TerminationReason.IterationTerminationCondition,
+                        terminationReason.toString(), scoreVsEpoch, bestModelEpoch, bestModelScore, epochCount,
+                        bestModel);
+                if (listener != null) {
+                    listener.onCompletion(result);
+                }
+                return result;
+            }
+
+            log.info("Completed training epoch {}", epochCount);
+
+
+            if ((epochCount == 0 && esConfig.getEvaluateEveryNEpochs() == 1)
+                    || epochCount % esConfig.getEvaluateEveryNEpochs() == 0) {
+                //Calculate score at this epoch:
+                ScoreCalculator sc = esConfig.getScoreCalculator();
+                double score = esConfig.getScoreCalculator().calculateScore(model);
+                scoreVsEpoch.put(epochCount, score);
+
+                boolean invalidScore = Double.isNaN(score) || Double.isInfinite(score);
+                if(invalidScore){
+                    log.warn("Score is not finite for epoch {}: score = {}", epochCount, score);
+                }
+
+                if ((sc.minimizeScore() && score < bestModelScore) || (!sc.minimizeScore() && score > bestModelScore) || (bestModelEpoch == -1 && invalidScore)) {
+                    //Save best model:
+                    if (bestModelEpoch == -1) {
+                        //First calculated/reported score
+                        log.info("Score at epoch {}: {}", epochCount, score);
+                    } else {
+                        log.info("New best model: score = {}, epoch = {} (previous: score = {}, epoch = {})", score,
+                                epochCount, bestModelScore, bestModelEpoch);
+                    }
+
+                    bestModelScore = score;
+                    bestModelEpoch = epochCount;
+
+                    try {
+                        esConfig.getModelSaver().saveBestModel(model, score);
+                    } catch (IOException e) {
+                        throw new RuntimeException("Error saving best model", e);
+                    }
+                } else {
+                    log.info("Score at epoch {}: {}", epochCount, score);
+                }
+
+                if (esConfig.isSaveLastModel()) {
+                    //Save last model:
+                    try {
+                        esConfig.getModelSaver().saveLatestModel(model, score);
+                    } catch (IOException e) {
+                        throw new RuntimeException("Error saving most recent model", e);
+                    }
+                }
+
+                if (listener != null) {
+                    listener.onEpoch(epochCount, score, esConfig, model);
+                }
+
+                //Check per-epoch termination conditions:
+                boolean epochTerminate = false;
+                EpochTerminationCondition termReason = null;
+                for (EpochTerminationCondition c : esConfig.getEpochTerminationConditions()) {
+                    if (c.terminate(epochCount, score, esConfig.getScoreCalculator().minimizeScore())) {
+                        epochTerminate = true;
+                        termReason = c;
+                        break;
+                    }
+                }
+                if (epochTerminate) {
+                    log.info("Hit epoch termination condition at epoch {}. Details: {}", epochCount,
+                            termReason);
+                    T bestModel;
+                    try {
+                        bestModel = esConfig.getModelSaver().getBestModel();
+                    } catch (IOException e2) {
+                        //Best model does not exist. Just save the current model
+                        if(esConfig.isSaveLastModel()) {
+                            try {
+                                esConfig.getModelSaver().saveBestModel(model,0.0);
+                                bestModel = model;
+                            } catch (IOException e) {
+                                log.error("Unable to save model.",e);
+                                throw new RuntimeException(e);
+                            }
+                        }
+                        else {
+                            log.error("Error with earlystopping",e2);
+                            throw new RuntimeException(e2);
+                        }
+
+                    }
+
+
+                    EarlyStoppingResult<T> result = new EarlyStoppingResult<>(
+                            EarlyStoppingResult.TerminationReason.EpochTerminationCondition,
+                            termReason.toString(), scoreVsEpoch, bestModelEpoch, bestModelScore, epochCount + 1,
+                            bestModel);
+                    if (listener != null) {
+                        listener.onCompletion(result);
+                    }
+
+                    return result;
+                }
+            }
+            epochCount++;
+
+        }
+    }
+
+    @Override
+    public void setListener(EarlyStoppingListener<T> listener) {
+        this.listener = listener;
+    }
+
+    //Trigger epoch listener methods manually - these won't be triggered due to not calling fit(DataSetIterator) etc
+    protected void triggerEpochListeners(boolean epochStart, Model model, int epochNum){
+        Collection<TrainingListener> listeners;
+        if(model instanceof MultiLayerNetwork){
+            MultiLayerNetwork n = ((MultiLayerNetwork) model);
+            listeners = n.getListeners();
+            n.setEpochCount(epochNum);
+        } else if(model instanceof ComputationGraph){
+            ComputationGraph cg = ((ComputationGraph) model);
+            listeners = cg.getListeners();
+            cg.getConfiguration().setEpochCount(epochNum);
+        } else {
+            return;
+        }
+
+        if(listeners != null && !listeners.isEmpty()){
+            for (TrainingListener l : listeners) {
+                if (epochStart) {
+                    l.onEpochStart(model);
+                } else {
+                    l.onEpochEnd(model);
+                }
+            }
+        }
+    }
+
+    protected void reset() {
+        if (train != null) {
+            train.reset();
+        }
+        if (trainMulti != null) {
+            trainMulti.reset();
+        }
+    }
+
+
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/earlystopping/trainer/EarlyStoppingGraphTrainer.java

@@ -0,0 +1,92 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.earlystopping.trainer;
+
+import org.deeplearning4j.datasets.iterator.utilty.SingletonDataSetIterator;
+import org.deeplearning4j.datasets.iterator.utilty.SingletonMultiDataSetIterator;
+import org.deeplearning4j.earlystopping.EarlyStoppingConfiguration;
+import org.deeplearning4j.earlystopping.listener.EarlyStoppingListener;
+import org.deeplearning4j.nn.graph.ComputationGraph;
+import org.nd4j.linalg.dataset.DataSet;
+import org.nd4j.linalg.dataset.api.MultiDataSet;
+import org.nd4j.linalg.dataset.api.iterator.DataSetIterator;
+import org.nd4j.linalg.dataset.api.iterator.MultiDataSetIterator;
+
+public class EarlyStoppingGraphTrainer extends BaseEarlyStoppingTrainer<ComputationGraph> { //implements IEarlyStoppingTrainer<ComputationGraph> {
+    private ComputationGraph net;
+
+    /**
+     * @param esConfig Configuration
+     * @param net Network to train using early stopping
+     * @param train DataSetIterator for training the network
+     */
+    public EarlyStoppingGraphTrainer(EarlyStoppingConfiguration<ComputationGraph> esConfig, ComputationGraph net,
+                    DataSetIterator train) {
+        this(esConfig, net, train, null);
+    }
+
+    /**Constructor for training using a {@link DataSetIterator}
+     * @param esConfig Configuration
+     * @param net Network to train using early stopping
+     * @param train DataSetIterator for training the network
+     * @param listener Early stopping listener. May be null.
+     */
+    public EarlyStoppingGraphTrainer(EarlyStoppingConfiguration<ComputationGraph> esConfig, ComputationGraph net,
+                    DataSetIterator train, EarlyStoppingListener<ComputationGraph> listener) {
+        super(esConfig, net, train, null, listener);
+        if (net.getNumInputArrays() != 1 || net.getNumOutputArrays() != 1)
+            throw new IllegalStateException(
+                            "Cannot do early stopping training on ComputationGraph with DataSetIterator: graph does not have 1 input and 1 output array");
+        this.net = net;
+    }
+
+    /**Constructor for training using a {@link MultiDataSetIterator}
+     * @param esConfig Configuration
+     * @param net Network to train using early stopping
+     * @param train DataSetIterator for training the network
+     * @param listener Early stopping listener. May be null.
+     */
+    public EarlyStoppingGraphTrainer(EarlyStoppingConfiguration<ComputationGraph> esConfig, ComputationGraph net,
+                    MultiDataSetIterator train, EarlyStoppingListener<ComputationGraph> listener) {
+        super(esConfig, net, null, train, listener);
+        this.net = net;
+    }
+
+    @Override
+    protected void fit(DataSet ds) {
+        net.fit(ds);
+    }
+
+    @Override
+    protected void fit(MultiDataSet mds) {
+        net.fit(mds);
+    }
+
+    @Override
+    protected void pretrain(DataSet ds) {
+        net.pretrain(new SingletonDataSetIterator(ds));
+    }
+
+    @Override
+    protected void pretrain(MultiDataSet mds) {
+        net.pretrain(new SingletonMultiDataSetIterator(mds));
+    }
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/earlystopping/trainer/EarlyStoppingTrainer.java

@@ -0,0 +1,76 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.earlystopping.trainer;
+
+import org.deeplearning4j.datasets.iterator.MultiDataSetWrapperIterator;
+import org.deeplearning4j.datasets.iterator.utilty.SingletonDataSetIterator;
+import org.deeplearning4j.datasets.iterator.utilty.SingletonMultiDataSetIterator;
+import org.deeplearning4j.earlystopping.EarlyStoppingConfiguration;
+import org.deeplearning4j.earlystopping.listener.EarlyStoppingListener;
+import org.deeplearning4j.nn.conf.MultiLayerConfiguration;
+import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;
+import org.nd4j.linalg.dataset.DataSet;
+import org.nd4j.linalg.dataset.api.MultiDataSet;
+import org.nd4j.linalg.dataset.api.iterator.DataSetIterator;
+
+public class EarlyStoppingTrainer extends BaseEarlyStoppingTrainer<MultiLayerNetwork> {
+
+    private MultiLayerNetwork net;
+    private boolean isMultiEpoch = false;
+
+
+    public EarlyStoppingTrainer(EarlyStoppingConfiguration<MultiLayerNetwork> earlyStoppingConfiguration,
+                    MultiLayerConfiguration configuration, DataSetIterator train) {
+        this(earlyStoppingConfiguration, new MultiLayerNetwork(configuration), train);
+        net.init();
+    }
+
+    public EarlyStoppingTrainer(EarlyStoppingConfiguration<MultiLayerNetwork> esConfig, MultiLayerNetwork net,
+                    DataSetIterator train) {
+        this(esConfig, net, train, null);
+    }
+
+    public EarlyStoppingTrainer(EarlyStoppingConfiguration<MultiLayerNetwork> esConfig, MultiLayerNetwork net,
+                    DataSetIterator train, EarlyStoppingListener<MultiLayerNetwork> listener) {
+        super(esConfig, net, train, null, listener);
+        this.net = net;
+    }
+
+    @Override
+    protected void fit(DataSet ds) {
+        net.fit(ds);
+    }
+
+    @Override
+    protected void fit(MultiDataSet mds) {
+        net.fit(mds);
+    }
+
+    @Override
+    protected void pretrain(DataSet ds) {
+        net.pretrain(new SingletonDataSetIterator(ds));
+    }
+
+    @Override
+    protected void pretrain(MultiDataSet mds) {
+        net.pretrain(new MultiDataSetWrapperIterator(new SingletonMultiDataSetIterator(mds)));
+    }
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/earlystopping/trainer/IEarlyStoppingTrainer.java

@@ -0,0 +1,37 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.earlystopping.trainer;
+
+import org.deeplearning4j.earlystopping.EarlyStoppingResult;
+import org.deeplearning4j.earlystopping.listener.EarlyStoppingListener;
+import org.deeplearning4j.nn.api.Model;
+
+public interface IEarlyStoppingTrainer<T extends Model> {
+
+    /** Conduct early stopping training */
+    EarlyStoppingResult<T> fit();
+
+    EarlyStoppingResult<T> pretrain();
+
+    /** Set the early stopping listener */
+    void setListener(EarlyStoppingListener<T> listener);
+
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/eval/BaseEvaluation.java

@@ -0,0 +1,68 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.eval;
+
+import lombok.EqualsAndHashCode;
+import lombok.Getter;
+import org.nd4j.common.primitives.AtomicBoolean;
+import org.nd4j.common.primitives.AtomicDouble;
+import org.nd4j.common.primitives.serde.JsonDeserializerAtomicBoolean;
+import org.nd4j.common.primitives.serde.JsonDeserializerAtomicDouble;
+import org.nd4j.common.primitives.serde.JsonSerializerAtomicBoolean;
+import org.nd4j.common.primitives.serde.JsonSerializerAtomicDouble;
+import org.nd4j.shade.jackson.annotation.JsonAutoDetect;
+import org.nd4j.shade.jackson.core.JsonProcessingException;
+import org.nd4j.shade.jackson.databind.DeserializationFeature;
+import org.nd4j.shade.jackson.databind.MapperFeature;
+import org.nd4j.shade.jackson.databind.ObjectMapper;
+import org.nd4j.shade.jackson.databind.SerializationFeature;
+import org.nd4j.shade.jackson.databind.module.SimpleModule;
+import org.nd4j.shade.jackson.dataformat.yaml.YAMLFactory;
+
+@Deprecated
+@EqualsAndHashCode
+public abstract class BaseEvaluation<T extends BaseEvaluation> extends org.nd4j.evaluation.BaseEvaluation<T> {
+
+    @Getter
+    private static ObjectMapper objectMapper = configureMapper(new ObjectMapper());
+    @Getter
+    private static ObjectMapper yamlMapper = configureMapper(new ObjectMapper(new YAMLFactory()));
+
+    private static ObjectMapper configureMapper(ObjectMapper ret) {
+        ret.configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false);
+        ret.configure(SerializationFeature.FAIL_ON_EMPTY_BEANS, false);
+        ret.configure(MapperFeature.SORT_PROPERTIES_ALPHABETICALLY, false);
+        ret.enable(SerializationFeature.INDENT_OUTPUT);
+        SimpleModule atomicModule = new SimpleModule();
+        atomicModule.addSerializer(AtomicDouble.class,new JsonSerializerAtomicDouble());
+        atomicModule.addSerializer(AtomicBoolean.class,new JsonSerializerAtomicBoolean());
+        atomicModule.addDeserializer(AtomicDouble.class,new JsonDeserializerAtomicDouble());
+        atomicModule.addDeserializer(AtomicBoolean.class,new JsonDeserializerAtomicBoolean());
+        ret.registerModule(atomicModule);
+        //Serialize fields only, not using getters
+        ret.setVisibilityChecker(ret.getSerializationConfig().getDefaultVisibilityChecker()
+                .withFieldVisibility(JsonAutoDetect.Visibility.ANY)
+                .withGetterVisibility(JsonAutoDetect.Visibility.NONE)
+                .withSetterVisibility(JsonAutoDetect.Visibility.NONE)
+                .withCreatorVisibility(JsonAutoDetect.Visibility.NONE));
+        return ret;
+    }
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/eval/ConfusionMatrix.java

@@ -0,0 +1,58 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.eval;
+
+import org.nd4j.shade.guava.collect.HashMultiset;
+import org.nd4j.shade.guava.collect.Multiset;
+
+import java.io.Serializable;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+
+@Deprecated
+public class ConfusionMatrix<T extends Comparable<? super T>> extends org.nd4j.evaluation.classification.ConfusionMatrix<T> {
+
+    /**
+     * @deprecated Use {@link org.nd4j.evaluation.classification.ConfusionMatrix}
+     */
+    @Deprecated
+    public ConfusionMatrix(List<T> classes) {
+        super(classes);
+    }
+
+    /**
+     * @deprecated Use {@link org.nd4j.evaluation.classification.ConfusionMatrix}
+     */
+    @Deprecated
+    public ConfusionMatrix() {
+        super();
+    }
+
+    /**
+     * @deprecated Use {@link org.nd4j.evaluation.classification.ConfusionMatrix}
+     */
+    @Deprecated
+    public ConfusionMatrix(ConfusionMatrix<T> other) {
+        super(other);
+    }
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/eval/Evaluation.java

@@ -0,0 +1,196 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.eval;
+
+import lombok.EqualsAndHashCode;
+import lombok.NonNull;
+import org.nd4j.evaluation.EvaluationAveraging;
+import org.nd4j.evaluation.IEvaluation;
+import org.nd4j.linalg.api.ndarray.INDArray;
+
+import java.util.List;
+import java.util.Map;
+
+@EqualsAndHashCode(callSuper = true)
+@Deprecated
+public class Evaluation extends org.nd4j.evaluation.classification.Evaluation implements org.deeplearning4j.eval.IEvaluation<org.nd4j.evaluation.classification.Evaluation> {
+
+    /**
+     * Use {@link org.nd4j.evaluation.classification.Evaluation.Metric}
+     */
+    @Deprecated
+    public enum Metric {ACCURACY, F1, PRECISION, RECALL, GMEASURE, MCC;
+        public org.nd4j.evaluation.classification.Evaluation.Metric toNd4j(){
+            switch (this){
+                case ACCURACY:
+                    return org.nd4j.evaluation.classification.Evaluation.Metric.ACCURACY;
+                case F1:
+                    return org.nd4j.evaluation.classification.Evaluation.Metric.F1;
+                case PRECISION:
+                    return org.nd4j.evaluation.classification.Evaluation.Metric.PRECISION;
+                case RECALL:
+                    return org.nd4j.evaluation.classification.Evaluation.Metric.RECALL;
+                case GMEASURE:
+                    return org.nd4j.evaluation.classification.Evaluation.Metric.GMEASURE;
+                case MCC:
+                    return org.nd4j.evaluation.classification.Evaluation.Metric.MCC;
+                default:
+                    throw new IllegalStateException("Unknown enum state: " + this);
+            }
+        }
+    }
+
+    /**
+     * @deprecated Use ND4J Evaluation class, which has the same interface: {@link org.nd4j.evaluation.classification.Evaluation.Metric}
+     */
+    @Deprecated
+    public Evaluation() {
+    }
+
+    /**
+     * @deprecated Use ND4J Evaluation class, which has the same interface: {@link org.nd4j.evaluation.classification.Evaluation.Metric}
+     */
+    @Deprecated
+    public Evaluation(int numClasses) {
+        super(numClasses);
+    }
+
+    /**
+     * @deprecated Use ND4J Evaluation class, which has the same interface: {@link org.nd4j.evaluation.classification.Evaluation.Metric}
+     */
+    @Deprecated
+    public Evaluation(int numClasses, Integer binaryPositiveClass){
+        super(numClasses, binaryPositiveClass);
+    }
+
+    /**
+     * @deprecated Use ND4J Evaluation class, which has the same interface: {@link org.nd4j.evaluation.classification.Evaluation.Metric}
+     */
+    @Deprecated
+    public Evaluation(List<String> labels) {
+        super(labels);
+    }
+
+    /**
+     * @deprecated Use ND4J Evaluation class, which has the same interface: {@link org.nd4j.evaluation.classification.Evaluation.Metric}
+     */
+    @Deprecated
+    public Evaluation(Map<Integer, String> labels) {
+        super(labels);
+    }
+
+    /**
+     * @deprecated Use ND4J Evaluation class, which has the same interface: {@link org.nd4j.evaluation.classification.Evaluation.Metric}
+     */
+    @Deprecated
+    public Evaluation(List<String> labels, int topN) {
+        super(labels, topN);
+    }
+
+    /**
+     * @deprecated Use ND4J Evaluation class, which has the same interface: {@link org.nd4j.evaluation.classification.Evaluation.Metric}
+     */
+    @Deprecated
+    public Evaluation(double binaryDecisionThreshold) {
+        super(binaryDecisionThreshold);
+    }
+
+    /**
+     * @deprecated Use ND4J Evaluation class, which has the same interface: {@link org.nd4j.evaluation.classification.Evaluation.Metric}
+     */
+    @Deprecated
+    public Evaluation(double binaryDecisionThreshold, @NonNull Integer binaryPositiveClass) {
+        super(binaryDecisionThreshold, binaryPositiveClass);
+    }
+
+    /**
+     * @deprecated Use ND4J Evaluation class, which has the same interface: {@link org.nd4j.evaluation.classification.Evaluation.Metric}
+     */
+    @Deprecated
+    public Evaluation(INDArray costArray) {
+        super(costArray);
+    }
+
+    /**
+     * @deprecated Use ND4J Evaluation class, which has the same interface: {@link org.nd4j.evaluation.classification.Evaluation.Metric}
+     */
+    @Deprecated
+    public Evaluation(List<String> labels, INDArray costArray) {
+        super(labels, costArray);
+    }
+
+    @Deprecated
+    public double precision(org.deeplearning4j.eval.EvaluationAveraging averaging) {
+        return precision(averaging.toNd4j());
+    }
+
+    @Deprecated
+    public double recall(org.deeplearning4j.eval.EvaluationAveraging averaging) {
+        return recall(averaging.toNd4j());
+    }
+
+    public double falsePositiveRate(org.deeplearning4j.eval.EvaluationAveraging averaging) {
+        return falsePositiveRate(averaging.toNd4j());
+    }
+
+    public double falseNegativeRate(org.deeplearning4j.eval.EvaluationAveraging averaging) {
+        return falseNegativeRate(averaging.toNd4j());
+    }
+
+    public double f1(org.deeplearning4j.eval.EvaluationAveraging averaging) {
+        return f1(averaging.toNd4j());
+    }
+
+    public double fBeta(double beta, org.deeplearning4j.eval.EvaluationAveraging averaging) {
+        return fBeta(beta, averaging.toNd4j());
+    }
+
+    public double gMeasure(org.deeplearning4j.eval.EvaluationAveraging averaging) {
+        return gMeasure(averaging.toNd4j());
+    }
+
+    public double matthewsCorrelation(org.deeplearning4j.eval.EvaluationAveraging averaging) {
+        return matthewsCorrelation(averaging.toNd4j());
+    }
+
+    /**
+     * @deprecated Use ND4J Evaluation class, which has the same interface: {@link org.nd4j.evaluation.classification.Evaluation.Metric}
+     */
+    public double scoreForMetric(Metric metric){
+        return scoreForMetric(metric.toNd4j());
+    }
+
+    /**
+     * @deprecated Use ND4J Evaluation class, which has the same interface: {@link org.nd4j.evaluation.classification.Evaluation.Metric}
+     */
+    @Deprecated
+    public static Evaluation fromJson(String json) {
+        return fromJson(json, Evaluation.class);
+    }
+
+    /**
+     * @deprecated Use ND4J Evaluation class, which has the same interface: {@link org.nd4j.evaluation.classification.Evaluation.Metric}
+     */
+    @Deprecated
+    public static Evaluation fromYaml(String yaml) {
+        return fromYaml(yaml, Evaluation.class);
+    }
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/eval/EvaluationAveraging.java

@@ -0,0 +1,36 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.eval;
+
+@Deprecated
+public enum EvaluationAveraging {
+    Macro, Micro;
+
+    public org.nd4j.evaluation.EvaluationAveraging toNd4j(){
+        switch (this){
+            case Macro:
+                return org.nd4j.evaluation.EvaluationAveraging.Macro;
+            case Micro:
+                return org.nd4j.evaluation.EvaluationAveraging.Micro;
+        }
+        throw new UnsupportedOperationException("Unknown: " + this);
+    }
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/eval/EvaluationBinary.java

@@ -0,0 +1,71 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.eval;
+
+import lombok.Data;
+import lombok.EqualsAndHashCode;
+import lombok.NoArgsConstructor;
+import org.nd4j.linalg.api.ndarray.INDArray;
+
+@Deprecated
+@NoArgsConstructor
+@EqualsAndHashCode(callSuper = true)
+@Data
+public class EvaluationBinary extends org.nd4j.evaluation.classification.EvaluationBinary implements IEvaluation<org.nd4j.evaluation.classification.EvaluationBinary> {
+    @Deprecated
+    public static final int DEFAULT_PRECISION = 4;
+    @Deprecated
+    public static final double DEFAULT_EDGE_VALUE = 0.0;
+
+    /**
+     * Use {@link org.nd4j.evaluation.classification.EvaluationBinary}
+     */
+    @Deprecated
+    public EvaluationBinary(INDArray decisionThreshold) {
+        super(decisionThreshold);
+    }
+
+    /**
+     * Use {@link org.nd4j.evaluation.classification.EvaluationBinary}
+     */
+    @Deprecated
+    public EvaluationBinary(int size, Integer rocBinarySteps) {
+        super(size, rocBinarySteps);
+    }
+
+    /**
+     * Use {@link org.nd4j.evaluation.classification.EvaluationBinary#fromJson(String)}
+     */
+    @Deprecated
+    public static EvaluationBinary fromJson(String json) {
+        return fromJson(json, EvaluationBinary.class);
+    }
+
+    /**
+     * Use {@link org.nd4j.evaluation.classification.EvaluationBinary.fromYaml(String)}
+     */
+    @Deprecated
+    public static EvaluationBinary fromYaml(String yaml) {
+        return fromYaml(yaml, EvaluationBinary.class);
+    }
+
+
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/eval/EvaluationCalibration.java

@@ -0,0 +1,57 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.eval;
+
+import lombok.EqualsAndHashCode;
+import lombok.Getter;
+import org.nd4j.shade.jackson.annotation.JsonProperty;
+
+@Deprecated
+@Getter
+@EqualsAndHashCode(callSuper = true)
+public class EvaluationCalibration extends org.nd4j.evaluation.classification.EvaluationCalibration implements IEvaluation<org.nd4j.evaluation.classification.EvaluationCalibration> {
+
+    /**
+     * @deprecated Use {@link org.nd4j.evaluation.classification.EvaluationCalibration}
+     */
+    @Deprecated
+    public EvaluationCalibration() {
+        super();
+    }
+
+    /**
+     * @deprecated Use {@link org.nd4j.evaluation.classification.EvaluationCalibration}
+     */
+    @Deprecated
+    public EvaluationCalibration(int reliabilityDiagNumBins, int histogramNumBins) {
+        super(reliabilityDiagNumBins, histogramNumBins);
+    }
+
+    /**
+     * @deprecated Use {@link org.nd4j.evaluation.classification.EvaluationCalibration}
+     */
+    @Deprecated
+    public EvaluationCalibration(@JsonProperty("reliabilityDiagNumBins") int reliabilityDiagNumBins,
+                    @JsonProperty("histogramNumBins") int histogramNumBins,
+                    @JsonProperty("excludeEmptyBins") boolean excludeEmptyBins) {
+        super(reliabilityDiagNumBins, histogramNumBins, excludeEmptyBins);
+    }
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/eval/EvaluationUtils.java

@@ -0,0 +1,38 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.eval;
+
+import org.nd4j.common.base.Preconditions;
+import org.nd4j.evaluation.IEvaluation;
+
+@Deprecated
+public class EvaluationUtils extends org.nd4j.evaluation.EvaluationUtils {
+
+
+    public static <T> T copyToLegacy(IEvaluation<?> from, Class<T> to){
+        if(from == null)
+            return null;
+        Preconditions.checkState(to.isAssignableFrom(from.getClass()), "Invalid classes: %s vs %s", from.getClass(), to);
+
+
+        throw new UnsupportedOperationException("Not implemented");
+    }
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/eval/IEvaluation.java

@@ -0,0 +1,29 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.eval;
+
+import org.nd4j.shade.jackson.annotation.JsonTypeInfo;
+
+@Deprecated
+@JsonTypeInfo(use = JsonTypeInfo.Id.CLASS, include = JsonTypeInfo.As.PROPERTY)
+public interface IEvaluation<T extends org.nd4j.evaluation.IEvaluation> extends org.nd4j.evaluation.IEvaluation<T> {
+
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/eval/ROC.java

@@ -0,0 +1,82 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.eval;
+
+import lombok.Data;
+import lombok.EqualsAndHashCode;
+import lombok.NoArgsConstructor;
+
+@Deprecated
+@EqualsAndHashCode(callSuper = true)
+@Data
+public class ROC extends org.nd4j.evaluation.classification.ROC implements IEvaluation<org.nd4j.evaluation.classification.ROC> {
+
+    /**
+     * @deprecated Use {@link org.nd4j.evaluation.classification.ROC}
+     */
+    @Deprecated
+    public ROC() { }
+
+    /**
+     * @deprecated Use {@link org.nd4j.evaluation.classification.ROC}
+     */
+    @Deprecated
+    public ROC(int thresholdSteps) {
+        super(thresholdSteps);
+    }
+
+    /**
+     * @deprecated Use {@link org.nd4j.evaluation.classification.ROC}
+     */
+    @Deprecated
+    public ROC(int thresholdSteps, boolean rocRemoveRedundantPts) {
+        super(thresholdSteps, rocRemoveRedundantPts);
+    }
+
+    /**
+     * @deprecated Use {@link org.nd4j.evaluation.classification.ROC}
+     */
+    @Deprecated
+    public ROC(int thresholdSteps, boolean rocRemoveRedundantPts, int exactAllocBlockSize) {
+        super(thresholdSteps, rocRemoveRedundantPts, exactAllocBlockSize);
+    }
+
+    /**
+     * @deprecated Use {@link org.nd4j.evaluation.classification.ROC.CountsForThreshold}
+     */
+    @Deprecated
+    @NoArgsConstructor
+    public static class CountsForThreshold extends org.nd4j.evaluation.classification.ROC.CountsForThreshold {
+
+        public CountsForThreshold(double threshold) {
+            super(threshold);
+        }
+
+        public CountsForThreshold(double threshold, long countTruePositive, long countFalsePositive){
+            super(threshold, countTruePositive, countFalsePositive);
+        }
+
+        @Override
+        public CountsForThreshold clone() {
+            return new CountsForThreshold(getThreshold(), getCountTruePositive(), getCountFalsePositive());
+        }
+    }
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/eval/ROCBinary.java

@@ -0,0 +1,55 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.eval;
+
+import lombok.EqualsAndHashCode;
+
+@Deprecated
+@EqualsAndHashCode(callSuper = true)
+public class ROCBinary extends org.nd4j.evaluation.classification.ROCBinary implements IEvaluation<org.nd4j.evaluation.classification.ROCBinary> {
+    /**
+     * @deprecated Use {@link org.nd4j.evaluation.classification.ROCBinary}
+     */
+    @Deprecated
+    public static final int DEFAULT_STATS_PRECISION = 4;
+
+    /**
+     * @deprecated Use {@link org.nd4j.evaluation.classification.ROCBinary}
+     */
+    @Deprecated
+    public ROCBinary() { }
+
+    /**
+     * @deprecated Use {@link org.nd4j.evaluation.classification.ROCBinary}
+     */
+    @Deprecated
+    public ROCBinary(int thresholdSteps) {
+        super(thresholdSteps);
+    }
+
+    /**
+     * @deprecated Use {@link org.nd4j.evaluation.classification.ROCBinary}
+     */
+    @Deprecated
+    public ROCBinary(int thresholdSteps, boolean rocRemoveRedundantPts) {
+        super(thresholdSteps, rocRemoveRedundantPts);
+    }
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/eval/ROCMultiClass.java

@@ -0,0 +1,57 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.eval;
+
+import lombok.Data;
+import lombok.EqualsAndHashCode;
+
+@Deprecated
+@Data
+@EqualsAndHashCode(callSuper = true)
+public class ROCMultiClass extends org.nd4j.evaluation.classification.ROCMultiClass implements IEvaluation<org.nd4j.evaluation.classification.ROCMultiClass> {
+    /**
+     * @deprecated Use {@link org.nd4j.evaluation.classification.ROCMultiClass}
+     */
+    @Deprecated
+    public static final int DEFAULT_STATS_PRECISION = 4;
+
+    /**
+     * @deprecated Use {@link org.nd4j.evaluation.classification.ROCMultiClass}
+     */
+    @Deprecated
+    public ROCMultiClass() { }
+
+    /**
+     * @deprecated Use {@link org.nd4j.evaluation.classification.ROCMultiClass}
+     */
+    @Deprecated
+    public ROCMultiClass(int thresholdSteps) {
+        super(thresholdSteps);
+    }
+
+    /**
+     * @deprecated Use {@link org.nd4j.evaluation.classification.ROCMultiClass}
+     */
+    @Deprecated
+    public ROCMultiClass(int thresholdSteps, boolean rocRemoveRedundantPts) {
+        super(thresholdSteps, rocRemoveRedundantPts);
+    }
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/eval/RegressionEvaluation.java

@@ -0,0 +1,115 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.eval;
+
+import lombok.Data;
+import lombok.EqualsAndHashCode;
+
+import java.util.List;
+
+@Deprecated
+@Data
+@EqualsAndHashCode(callSuper = true)
+public class RegressionEvaluation extends org.nd4j.evaluation.regression.RegressionEvaluation implements IEvaluation<org.nd4j.evaluation.regression.RegressionEvaluation> {
+
+    /**
+     * @deprecated Use ND4J RegressionEvaluation class, which has the same interface: {@link org.nd4j.evaluation.regression.RegressionEvaluation.Metric}
+     */
+    @Deprecated
+    public enum Metric { MSE, MAE, RMSE, RSE, PC, R2;
+        public boolean minimize(){
+            return toNd4j().minimize();
+        }
+
+        public org.nd4j.evaluation.regression.RegressionEvaluation.Metric toNd4j(){
+            switch (this){
+                case MSE:
+                    return org.nd4j.evaluation.regression.RegressionEvaluation.Metric.MSE;
+                case MAE:
+                    return org.nd4j.evaluation.regression.RegressionEvaluation.Metric.MAE;
+                case RMSE:
+                    return org.nd4j.evaluation.regression.RegressionEvaluation.Metric.RMSE;
+                case RSE:
+                    return org.nd4j.evaluation.regression.RegressionEvaluation.Metric.RSE;
+                case PC:
+                    return org.nd4j.evaluation.regression.RegressionEvaluation.Metric.PC;
+                case R2:
+                    return org.nd4j.evaluation.regression.RegressionEvaluation.Metric.R2;
+                default:
+                    throw new IllegalStateException("Unknown enum: " + this);
+            }
+        }
+    }
+
+    /**
+     * @deprecated Use ND4J RegressionEvaluation class, which has the same interface: {@link org.nd4j.evaluation.regression.RegressionEvaluation}
+     */
+    @Deprecated
+    public RegressionEvaluation() { }
+
+    /**
+     * @deprecated Use ND4J RegressionEvaluation class, which has the same interface: {@link org.nd4j.evaluation.regression.RegressionEvaluation}
+     */
+    @Deprecated
+    public RegressionEvaluation(long nColumns) {
+        super(nColumns);
+    }
+
+    /**
+     * @deprecated Use ND4J RegressionEvaluation class, which has the same interface: {@link org.nd4j.evaluation.regression.RegressionEvaluation}
+     */
+    @Deprecated
+    public RegressionEvaluation(long nColumns, long precision) {
+        super(nColumns, precision);
+    }
+
+    /**
+     * @deprecated Use ND4J RegressionEvaluation class, which has the same interface: {@link org.nd4j.evaluation.regression.RegressionEvaluation}
+     */
+    @Deprecated
+    public RegressionEvaluation(String... columnNames) {
+        super(columnNames);
+    }
+
+    /**
+     * @deprecated Use ND4J RegressionEvaluation class, which has the same interface: {@link org.nd4j.evaluation.regression.RegressionEvaluation}
+     */
+    @Deprecated
+    public RegressionEvaluation(List<String> columnNames) {
+        super(columnNames);
+    }
+
+    /**
+     * @deprecated Use ND4J RegressionEvaluation class, which has the same interface: {@link org.nd4j.evaluation.regression.RegressionEvaluation}
+     */
+    @Deprecated
+    public RegressionEvaluation(List<String> columnNames, long precision) {
+        super(columnNames, precision);
+    }
+
+    /**
+     * @deprecated Use ND4J RegressionEvaluation class, which has the same interface: {@link org.nd4j.evaluation.regression.RegressionEvaluation}
+     */
+    @Deprecated
+    public double scoreForMetric(Metric metric){
+        return scoreForMetric(metric.toNd4j());
+    }
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/eval/curves/Histogram.java

@@ -0,0 +1,52 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.eval.curves;
+
+import lombok.Data;
+import org.nd4j.evaluation.curves.BaseHistogram;
+import org.nd4j.shade.jackson.annotation.JsonProperty;
+
+@Deprecated
+@Data
+public class Histogram extends org.nd4j.evaluation.curves.Histogram {
+
+    /**
+     * @deprecated Use {@link org.nd4j.evaluation.curves.Histogram}
+     */
+    public Histogram(@JsonProperty("title") String title, @JsonProperty("lower") double lower,
+                    @JsonProperty("upper") double upper, @JsonProperty("binCounts") int[] binCounts) {
+        super(title, lower, upper, binCounts);
+    }
+
+    /**
+     * @deprecated Use {@link org.nd4j.evaluation.curves.Histogram}
+     */
+    public static Histogram fromJson(String json) {
+        return BaseHistogram.fromJson(json, Histogram.class);
+    }
+
+    /**
+     * @deprecated Use {@link org.nd4j.evaluation.curves.Histogram}
+     */
+    public static Histogram fromYaml(String yaml) {
+        return BaseHistogram.fromYaml(yaml, Histogram.class);
+    }
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/eval/curves/PrecisionRecallCurve.java

@@ -0,0 +1,57 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.eval.curves;
+
+import org.nd4j.shade.guava.base.Preconditions;
+import lombok.Data;
+import lombok.EqualsAndHashCode;
+import org.nd4j.shade.jackson.annotation.JsonProperty;
+
+import java.util.Arrays;
+
+@Deprecated
+@Data
+@EqualsAndHashCode(callSuper = true)
+public class PrecisionRecallCurve extends org.nd4j.evaluation.curves.PrecisionRecallCurve{
+
+    /**
+     * @deprecated Use {@link org.nd4j.evaluation.curves.ReliabilityDiagram}
+     */
+    @Deprecated
+    public PrecisionRecallCurve(@JsonProperty("threshold") double[] threshold,
+                    @JsonProperty("precision") double[] precision, @JsonProperty("recall") double[] recall,
+                    @JsonProperty("tpCount") int[] tpCount, @JsonProperty("fpCount") int[] fpCount,
+                    @JsonProperty("fnCount") int[] fnCount, @JsonProperty("totalCount") int totalCount) {
+        super(threshold, precision, recall, tpCount, fpCount, fnCount, totalCount);
+    }
+
+    public static class Point extends org.nd4j.evaluation.curves.PrecisionRecallCurve.Point{
+        public Point(int idx, double threshold, double precision, double recall) {
+            super(idx, threshold, precision, recall);
+        }
+    }
+
+    public static class Confusion extends org.nd4j.evaluation.curves.PrecisionRecallCurve.Confusion{
+        public Confusion(org.nd4j.evaluation.curves.PrecisionRecallCurve.Point point, int tpCount, int fpCount, int fnCount, int tnCount) {
+            super(point, tpCount, fpCount, fnCount, tnCount);
+        }
+    }
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/eval/curves/ReliabilityDiagram.java

@@ -0,0 +1,38 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.eval.curves;
+
+import lombok.NonNull;
+import org.nd4j.shade.jackson.annotation.JsonProperty;
+
+@Deprecated
+public class ReliabilityDiagram extends org.nd4j.evaluation.curves.ReliabilityDiagram {
+
+    /**
+     * @deprecated Use {@link org.nd4j.evaluation.curves.ReliabilityDiagram}
+     */
+    @Deprecated
+    public ReliabilityDiagram(@JsonProperty("title") String title,
+                    @NonNull @JsonProperty("meanPredictedValueX") double[] meanPredictedValueX,
+                    @NonNull @JsonProperty("fractionPositivesY") double[] fractionPositivesY) {
+        super(title, meanPredictedValueX, fractionPositivesY);
+    }
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/eval/curves/RocCurve.java

@@ -0,0 +1,59 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.eval.curves;
+
+import org.nd4j.shade.guava.base.Preconditions;
+import lombok.Data;
+import lombok.EqualsAndHashCode;
+import org.nd4j.shade.jackson.annotation.JsonProperty;
+
+@Deprecated
+@Data
+@EqualsAndHashCode(exclude = {"auc"}, callSuper = false)
+public class RocCurve extends org.nd4j.evaluation.curves.RocCurve {
+
+    /**
+     * @deprecated Use {@link org.nd4j.evaluation.curves.RocCurve}
+     */
+    @Deprecated
+    public RocCurve(@JsonProperty("threshold") double[] threshold, @JsonProperty("fpr") double[] fpr,
+                    @JsonProperty("tpr") double[] tpr) {
+        super(threshold, fpr, tpr);
+    }
+
+
+    /**
+     * @deprecated Use {@link org.nd4j.evaluation.curves.RocCurve}
+     */
+    @Deprecated
+    public static RocCurve fromJson(String json) {
+        return fromJson(json, RocCurve.class);
+    }
+
+    /**
+     * @deprecated Use {@link org.nd4j.evaluation.curves.RocCurve}
+     */
+    @Deprecated
+    public static RocCurve fromYaml(String yaml) {
+        return fromYaml(yaml, RocCurve.class);
+    }
+
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/eval/meta/Prediction.java

@@ -0,0 +1,30 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.eval.meta;
+
+import lombok.Data;
+
+@Data
+public class Prediction extends org.nd4j.evaluation.meta.Prediction {
+    public Prediction(int actualClass, int predictedClass, Object recordMetaData){
+        super(actualClass, predictedClass, recordMetaData);
+    }
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/exception/DL4JException.java

@@ -0,0 +1,38 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.exception;
+
+public class DL4JException extends RuntimeException {
+
+    public DL4JException() {}
+
+    public DL4JException(String message) {
+        super(message);
+    }
+
+    public DL4JException(String message, Throwable cause) {
+        super(message, cause);
+    }
+
+    public DL4JException(Throwable cause) {
+        super(cause);
+    }
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/exception/DL4JInvalidConfigException.java

@@ -0,0 +1,37 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.exception;
+
+public class DL4JInvalidConfigException extends DL4JException {
+    public DL4JInvalidConfigException() {}
+
+    public DL4JInvalidConfigException(String message) {
+        super(message);
+    }
+
+    public DL4JInvalidConfigException(String message, Throwable cause) {
+        super(message, cause);
+    }
+
+    public DL4JInvalidConfigException(Throwable cause) {
+        super(cause);
+    }
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/exception/DL4JInvalidInputException.java

@@ -0,0 +1,38 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.exception;
+
+public class DL4JInvalidInputException extends DL4JException {
+
+    public DL4JInvalidInputException() {}
+
+    public DL4JInvalidInputException(String message) {
+        super(message);
+    }
+
+    public DL4JInvalidInputException(String message, Throwable cause) {
+        super(message, cause);
+    }
+
+    public DL4JInvalidInputException(Throwable cause) {
+        super(cause);
+    }
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/exception/DeepLearningException.java

@@ -0,0 +1,53 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.exception;
+
+public class DeepLearningException extends Exception {
+
+    /**
+     * 
+     */
+    private static final long serialVersionUID = -7973589163269627293L;
+
+    public DeepLearningException() {
+        super();
+    }
+
+    public DeepLearningException(String message, Throwable cause, boolean enableSuppression,
+                    boolean writableStackTrace) {
+        super(message, cause, enableSuppression, writableStackTrace);
+    }
+
+    public DeepLearningException(String message, Throwable cause) {
+        super(message, cause);
+    }
+
+    public DeepLearningException(String message) {
+        super(message);
+    }
+
+    public DeepLearningException(Throwable cause) {
+        super(cause);
+    }
+
+
+
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/exception/InvalidStepException.java

@@ -0,0 +1,91 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.exception;
+
+public class InvalidStepException extends Exception {
+
+    /**
+     * Constructs a new exception with the specified detail message.  The
+     * cause is not initialized, and may subsequently be initialized by
+     * a call to {@link #initCause}.
+     *
+     * @param message the detail message. The detail message is saved for
+     *                later retrieval by the {@link #getMessage()} method.
+     */
+    public InvalidStepException(String message) {
+        super(message);
+    }
+
+    /**
+     * Constructs a new exception with the specified detail message and
+     * cause.  <p>Note that the detail message associated with
+     * {@code cause} is <i>not</i> automatically incorporated in
+     * this exception's detail message.
+     *
+     * @param message the detail message (which is saved for later retrieval
+     *                by the {@link #getMessage()} method).
+     * @param cause   the cause (which is saved for later retrieval by the
+     *                {@link #getCause()} method).  (A <tt>null</tt> value is
+     *                permitted, and indicates that the cause is nonexistent or
+     *                unknown.)
+     * @since 1.4
+     */
+    public InvalidStepException(String message, Throwable cause) {
+        super(message, cause);
+    }
+
+    /**
+     * Constructs a new exception with the specified cause and a detail
+     * message of <tt>(cause==null ? null : cause.toString())</tt> (which
+     * typically contains the class and detail message of <tt>cause</tt>).
+     * This constructor is useful for exceptions that are little more than
+     * wrappers for other throwables (for example, {@link
+     * java.security.PrivilegedActionException}).
+     *
+     * @param cause the cause (which is saved for later retrieval by the
+     *              {@link #getCause()} method).  (A <tt>null</tt> value is
+     *              permitted, and indicates that the cause is nonexistent or
+     *              unknown.)
+     * @since 1.4
+     */
+    public InvalidStepException(Throwable cause) {
+        super(cause);
+    }
+
+    /**
+     * Constructs a new exception with the specified detail message,
+     * cause, suppression enabled or disabled, and writable stack
+     * trace enabled or disabled.
+     *
+     * @param message            the detail message.
+     * @param cause              the cause.  (A {@code null} value is permitted,
+     *                           and indicates that the cause is nonexistent or unknown.)
+     * @param enableSuppression  whether or not suppression is enabled
+     *                           or disabled
+     * @param writableStackTrace whether or not the stack trace should
+     *                           be writable
+     * @since 1.7
+     */
+    protected InvalidStepException(String message, Throwable cause, boolean enableSuppression,
+                    boolean writableStackTrace) {
+        super(message, cause, enableSuppression, writableStackTrace);
+    }
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/gradientcheck/GradientCheckUtil.java

@@ -0,0 +1,714 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.gradientcheck;
+
+import lombok.Data;
+import lombok.NoArgsConstructor;
+import lombok.experimental.Accessors;
+import lombok.extern.slf4j.Slf4j;
+import lombok.val;
+
+import org.nd4j.linalg.api.buffer.DataType;
+import org.nd4j.linalg.exception.ND4JArraySizeException;
+import org.nd4j.common.function.Consumer;
+import org.nd4j.linalg.lossfunctions.impl.LossBinaryXENT;
+import org.nd4j.common.primitives.Pair;
+import org.deeplearning4j.nn.api.Layer;
+import org.deeplearning4j.nn.api.Updater;
+import org.deeplearning4j.nn.api.layers.IOutputLayer;
+import org.deeplearning4j.nn.conf.NeuralNetConfiguration;
+import org.deeplearning4j.nn.conf.graph.GraphVertex;
+import org.deeplearning4j.nn.conf.graph.LayerVertex;
+import org.deeplearning4j.nn.conf.layers.BaseLayer;
+import org.deeplearning4j.nn.gradient.Gradient;
+import org.deeplearning4j.nn.graph.ComputationGraph;
+import org.deeplearning4j.nn.layers.BaseOutputLayer;
+import org.deeplearning4j.nn.layers.LossLayer;
+import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;
+import org.deeplearning4j.nn.updater.graph.ComputationGraphUpdater;
+import org.nd4j.linalg.activations.IActivation;
+import org.nd4j.linalg.activations.impl.ActivationSoftmax;
+import org.nd4j.linalg.api.buffer.util.DataTypeUtil;
+import org.nd4j.linalg.api.ndarray.INDArray;
+import org.nd4j.linalg.dataset.DataSet;
+import org.nd4j.linalg.dataset.MultiDataSet;
+import org.nd4j.linalg.factory.Nd4j;
+import org.nd4j.linalg.learning.config.IUpdater;
+import org.nd4j.linalg.learning.config.NoOp;
+import org.nd4j.linalg.learning.config.Sgd;
+import org.nd4j.linalg.lossfunctions.ILossFunction;
+import org.nd4j.linalg.lossfunctions.impl.LossMCXENT;
+import org.deeplearning4j.nn.workspace.LayerWorkspaceMgr;
+
+import java.util.*;
+
+@Slf4j
+public class GradientCheckUtil {
+
+
+    private GradientCheckUtil() {}
+
+    static {
+        Nd4j.setDefaultDataTypes(DataType.DOUBLE, DataType.DOUBLE);
+    }
+
+
+    private static void configureLossFnClippingIfPresent(IOutputLayer outputLayer) {
+
+        ILossFunction lfn = null;
+        IActivation afn = null;
+        if(outputLayer instanceof BaseOutputLayer) {
+            BaseOutputLayer o = (BaseOutputLayer)outputLayer;
+            lfn = ((org.deeplearning4j.nn.conf.layers.BaseOutputLayer)o.layerConf()).getLossFn();
+            afn = o.layerConf().getActivationFn();
+        } else if(outputLayer instanceof LossLayer){
+            LossLayer o = (LossLayer) outputLayer;
+            lfn = o.layerConf().getLossFn();
+            afn = o.layerConf().getActivationFn();
+        }
+
+        if (lfn instanceof LossMCXENT && afn instanceof ActivationSoftmax && ((LossMCXENT) lfn).getSoftmaxClipEps() != 0) {
+            log.info("Setting softmax clipping epsilon to 0.0 for " + lfn.getClass()
+                    + " loss function to avoid spurious gradient check failures");
+            ((LossMCXENT) lfn).setSoftmaxClipEps(0.0);
+        } else if(lfn instanceof LossBinaryXENT && ((LossBinaryXENT) lfn).getClipEps() != 0) {
+            log.info("Setting clipping epsilon to 0.0 for " + lfn.getClass()
+                    + " loss function to avoid spurious gradient check failures");
+            ((LossBinaryXENT) lfn).setClipEps(0.0);
+        }
+
+        log.info("Done setting clipping");
+    }
+
+    public enum PrintMode {
+        ALL,
+        ZEROS,
+        FAILURES_ONLY
+    }
+
+    @Accessors(fluent = true)
+    @Data
+    @NoArgsConstructor
+    public static class MLNConfig {
+        private MultiLayerNetwork net;
+        private INDArray input;
+        private INDArray labels;
+        private INDArray inputMask;
+        private INDArray labelMask;
+        private double epsilon = 1e-6;
+        private double maxRelError = 1e-3;
+        private double minAbsoluteError = 1e-8;
+        private PrintMode print = PrintMode.ZEROS;
+        private boolean exitOnFirstError = false;
+        private boolean subset;
+        private int maxPerParam;
+        private Set<String> excludeParams;
+        private Consumer<MultiLayerNetwork> callEachIter;
+    }
+
+    @Accessors(fluent = true)
+    @Data
+    @NoArgsConstructor
+    public static class GraphConfig {
+        private ComputationGraph net;
+        private INDArray[] inputs;
+        private INDArray[] labels;
+        private INDArray[] inputMask;
+        private INDArray[] labelMask;
+        private double epsilon = 1e-6;
+        private double maxRelError = 1e-3;
+        private double minAbsoluteError = 1e-8;
+        private PrintMode print = PrintMode.ZEROS;
+        private boolean exitOnFirstError = false;
+        private boolean subset;
+        private int maxPerParam;
+        private Set<String> excludeParams;
+        private Consumer<ComputationGraph> callEachIter;
+    }
+
+    /**
+     * Check backprop gradients for a MultiLayerNetwork.
+     * @param mln MultiLayerNetwork to test. This must be initialized.
+     * @param epsilon Usually on the order/ of 1e-4 or so.
+     * @param maxRelError Maximum relative error. Usually < 1e-5 or so, though maybe more for deep networks or those with nonlinear activation
+     * @param minAbsoluteError Minimum absolute error to cause a failure. Numerical gradients can be non-zero due to precision issues.
+     *                         For example, 0.0 vs. 1e-18: relative error is 1.0, but not really a failure
+     * @param print Whether to print full pass/failure details for each parameter gradient
+     * @param exitOnFirstError If true: return upon first failure. If false: continue checking even if
+     *  one parameter gradient has failed. Typically use false for debugging, true for unit tests.
+     * @param input Input array to use for forward pass. May be mini-batch data.
+     * @param labels Labels/targets to use to calculate backprop gradient. May be mini-batch data.
+     * @return true if gradients are passed, false otherwise.
+     */
+    @Deprecated
+    public static boolean checkGradients(MultiLayerNetwork mln, double epsilon, double maxRelError,
+                                         double minAbsoluteError, boolean print, boolean exitOnFirstError, INDArray input, INDArray labels) {
+        return checkGradients(new MLNConfig().net(mln)
+                .epsilon(epsilon)
+                .maxRelError(maxRelError)
+                .minAbsoluteError(minAbsoluteError)
+                .print(PrintMode.FAILURES_ONLY)
+                .exitOnFirstError(exitOnFirstError)
+                .input(input)
+                .labels(labels));
+    }
+
+    @Deprecated
+    public static boolean checkGradients(MultiLayerNetwork mln, double epsilon, double maxRelError,
+                                         double minAbsoluteError, boolean print, boolean exitOnFirstError,
+                                         INDArray input, INDArray labels, INDArray inputMask, INDArray labelMask,
+                                         boolean subset, int maxPerParam, Set<String> excludeParams, final Integer rngSeedResetEachIter) {
+        Consumer<MultiLayerNetwork> c = null;
+        if(rngSeedResetEachIter != null) {
+            c = multiLayerNetwork -> Nd4j.getRandom().setSeed(rngSeedResetEachIter);
+        }
+
+        return checkGradients(new MLNConfig().net(mln).epsilon(epsilon).maxRelError(maxRelError).minAbsoluteError(minAbsoluteError).print(PrintMode.FAILURES_ONLY)
+                .exitOnFirstError(exitOnFirstError).input(input).labels(labels).inputMask(inputMask).labelMask(labelMask).subset(subset).maxPerParam(maxPerParam).excludeParams(excludeParams).callEachIter(c));
+    }
+
+    public static boolean checkGradients(MLNConfig c) {
+        //Basic sanity checks on input:
+        if (c.epsilon <= 0.0 || c.epsilon > 0.1)
+            throw new IllegalArgumentException("Invalid epsilon: expect epsilon in range (0,0.1], usually 1e-4 or so");
+        if (c.maxRelError <= 0.0 || c.maxRelError > 0.25)
+            throw new IllegalArgumentException("Invalid maxRelativeError: " + c.maxRelError);
+        if (!(c.net.getOutputLayer() instanceof IOutputLayer))
+            throw new IllegalArgumentException("Cannot check backprop gradients without OutputLayer");
+
+        DataType dataType = DataTypeUtil.getDtypeFromContext();
+        if (dataType != DataType.DOUBLE) {
+            throw new IllegalStateException("Cannot perform gradient check: Datatype is not set to double precision ("
+                    + "is: " + dataType + "). Double precision must be used for gradient checks. Set "
+                    + "DataTypeUtil.setDTypeForContext(DataType.DOUBLE); before using GradientCheckUtil");
+        }
+
+        DataType netDataType = c.net.getLayerWiseConfigurations().getDataType();
+        if (netDataType != DataType.DOUBLE) {
+            throw new IllegalStateException("Cannot perform gradient check: Network datatype is not set to double precision ("
+                    + "is: " + netDataType + "). Double precision must be used for gradient checks. Create network with .dataType(DataType.DOUBLE) before using GradientCheckUtil");
+        }
+
+        if(netDataType != c.net.params().dataType()) {
+            throw new IllegalStateException("Parameters datatype does not match network configuration datatype ("
+                    + "is: " + c.net.params().dataType() + "). If network datatype is set to DOUBLE, parameters must also be DOUBLE.");
+        }
+
+
+        //Check network configuration:
+        int layerCount = 0;
+        for (NeuralNetConfiguration n : c.net.getLayerWiseConfigurations().getConfs()) {
+            if (n.getLayer() instanceof BaseLayer) {
+                BaseLayer bl = (BaseLayer) n.getLayer();
+                IUpdater u = bl.getIUpdater();
+                if (u instanceof Sgd) {
+                    //Must have LR of 1.0
+                    double lr = ((Sgd) u).getLearningRate();
+                    if (lr != 1.0) {
+                        throw new IllegalStateException("When using SGD updater, must also use lr=1.0 for layer "
+                                + layerCount + "; got " + u + " with lr=" + lr + " for layer \""
+                                + n.getLayer().getLayerName() + "\"");
+                    }
+                } else if (!(u instanceof NoOp)) {
+                    throw new IllegalStateException(
+                            "Must have Updater.NONE (or SGD + lr=1.0) for layer " + layerCount + "; got " + u);
+                }
+
+
+            }
+
+            if (n.getLayer().getIDropout() != null && c.callEachIter == null) {
+                throw new IllegalStateException("When gradient checking dropout, need to reset RNG seed each iter, or no" +
+                        " dropout should be present during gradient checks - got dropout = "
+                        + n.getLayer().getIDropout() + " for layer " + layerCount);
+            }
+        }
+
+        //Set softmax clipping to 0 if necessary, to avoid spurious failures due to clipping
+        for(Layer l : c.net.getLayers()) {
+            if(l instanceof IOutputLayer) {
+                configureLossFnClippingIfPresent((IOutputLayer) l);
+            }
+        }
+
+        c.net.setInput(c.input);
+        c.net.setLabels(c.labels);
+        c.net.setLayerMaskArrays(c.inputMask, c.labelMask);
+        if(c.callEachIter != null) {
+            c.callEachIter.accept(c.net);
+        }
+        c.net.computeGradientAndScore();
+        Pair<Gradient, Double> gradAndScore = c.net.gradientAndScore();
+
+        Updater updater = c.net().createUpdater();
+        updater.update(c.net, gradAndScore.getFirst(), 0, 0, c.net.batchSize(), LayerWorkspaceMgr.noWorkspaces());
+
+        INDArray gradientToCheck = gradAndScore.getFirst().gradient().dup(); //need dup: gradients are a *view* of the full gradient array (which will change every time backprop is done)
+        INDArray originalParams = c.net.params().dup(); //need dup: params are a *view* of full parameters
+
+        val nParams = originalParams.length();
+
+        Map<String, INDArray> paramTable = c.net.paramTable();
+        List<String> paramNames = new ArrayList<>(paramTable.keySet());
+        val paramEnds = new long[paramNames.size()];
+        paramEnds[0] = paramTable.get(paramNames.get(0)).length();
+        Map<String,Integer> stepSizeForParam;
+        if(c.subset) {
+            stepSizeForParam = new HashMap<>();
+            stepSizeForParam.put(paramNames.get(0), (int) Math.max(1, paramTable.get(paramNames.get(0)).length() / c.maxPerParam));
+        } else {
+            stepSizeForParam = null;
+        }
+        for (int i = 1; i < paramEnds.length; i++) {
+            val n = paramTable.get(paramNames.get(i)).length();
+            paramEnds[i] = paramEnds[i - 1] + n;
+            if(c.subset) {
+                long ss = n / c.maxPerParam;
+                if(ss == 0) {
+                    ss = n;
+                }
+
+                if (ss > Integer.MAX_VALUE)
+                    throw new ND4JArraySizeException();
+                stepSizeForParam.put(paramNames.get(i), (int) ss);
+            }
+        }
+
+        if(c.print == PrintMode.ALL) {
+            int i = 0;
+            for (Layer l : c.net.getLayers()) {
+                Set<String> s = l.paramTable().keySet();
+                log.info("Layer " + i + ": " + l.getClass().getSimpleName() + " - params " + s);
+                i++;
+            }
+        }
+
+
+        int totalNFailures = 0;
+        double maxError = 0.0;
+        DataSet ds = new DataSet(c.input, c.labels, c.inputMask, c.labelMask);
+        int currParamNameIdx = 0;
+
+        if(c.excludeParams != null && !c.excludeParams.isEmpty()) {
+            log.info("NOTE: parameters will be skipped due to config: {}", c.excludeParams);
+        }
+
+        INDArray params = c.net.params(); //Assumption here: params is a view that we can modify in-place
+        for (long i = 0; i < nParams;) {
+            //Get param name
+            if (i >= paramEnds[currParamNameIdx]) {
+                currParamNameIdx++;
+            }
+            String paramName = paramNames.get(currParamNameIdx);
+            if(c.excludeParams != null && c.excludeParams.contains(paramName)) {
+                i = paramEnds[currParamNameIdx++];
+                continue;
+            }
+
+            //(w+epsilon): Do forward pass and score
+            double origValue = params.getDouble(i);
+            params.putScalar(i, origValue + c.epsilon);
+            if(c.callEachIter != null) {
+                c.callEachIter.accept(c.net);
+            }
+            double scorePlus = c.net.score(ds, true);
+
+            //(w-epsilon): Do forward pass and score
+            params.putScalar(i, origValue - c.epsilon);
+            if(c.callEachIter != null) {
+                c.callEachIter.accept(c.net);
+            }
+            double scoreMinus = c.net.score(ds, true);
+
+            //Reset original param value
+            params.putScalar(i, origValue);
+
+            //Calculate numerical parameter gradient:
+            double scoreDelta = scorePlus - scoreMinus;
+
+            double numericalGradient = scoreDelta / (2 * c.epsilon);
+            if (Double.isNaN(numericalGradient))
+                throw new IllegalStateException("Numerical gradient was NaN for parameter " + i + " of " + nParams);
+
+            double backpropGradient = gradientToCheck.getDouble(i);
+            //http://cs231n.github.io/neural-networks-3/#gradcheck
+            //use mean centered
+            double relError = Math.abs(backpropGradient - numericalGradient)
+                    / (Math.abs(numericalGradient) + Math.abs(backpropGradient));
+            if (backpropGradient == 0.0 && numericalGradient == 0.0)
+                relError = 0.0; //Edge case: i.e., RNNs with time series length of 1.0
+
+            if (relError > maxError)
+                maxError = relError;
+            if (relError > c.maxRelError || Double.isNaN(relError)) {
+                double absError = Math.abs(backpropGradient - numericalGradient);
+                if (absError < c.minAbsoluteError) {
+                    if(c.print == PrintMode.ALL || c.print == PrintMode.ZEROS && absError == 0.0) {
+                        log.info("MLN Param " + i + " (" + paramName + ") passed: grad= " + backpropGradient
+                                + ", numericalGrad= " + numericalGradient + ", relError= " + relError
+                                + "; absolute error = " + absError + " < minAbsoluteError = " + c.minAbsoluteError);
+                    }
+                } else {
+                    log.info("MLN Param " + i + " (" + paramName + ") FAILED: grad= " + backpropGradient
+                            + ", numericalGrad= " + numericalGradient + ", relError= " + relError
+                            + ", scorePlus=" + scorePlus + ", scoreMinus= " + scoreMinus + ", paramValue = " + origValue);
+                    if (c.exitOnFirstError)
+                        return false;
+                    totalNFailures++;
+                }
+            } else if (c.print == PrintMode.ALL) {
+                log.info("Param " + i + " (" + paramName + ") passed: grad= " + backpropGradient + ", numericalGrad= "
+                        + numericalGradient + ", relError= " + relError);
+            }
+
+            long step;
+            if(c.subset) {
+                step = stepSizeForParam.get(paramName);
+                if(i + step > paramEnds[currParamNameIdx] + 1) {
+                    step = paramEnds[currParamNameIdx]+1 - i;
+                }
+            } else {
+                step = 1;
+            }
+
+            i += step;
+        }
+
+        val nPass = nParams - totalNFailures;
+        log.info("GradientCheckUtil.checkGradients(): " + nParams + " params checked, " + nPass + " passed, "
+                + totalNFailures + " failed. Largest relative error = " + maxError);
+
+        return totalNFailures == 0;
+    }
+
+    public static boolean checkGradients(GraphConfig c) {
+        //Basic sanity checks on input:
+        if (c.epsilon <= 0.0 || c.epsilon > 0.1)
+            throw new IllegalArgumentException("Invalid epsilon: expect epsilon in range (0,0.1], usually 1e-4 or so");
+        if (c.maxRelError <= 0.0 || c.maxRelError > 0.25)
+            throw new IllegalArgumentException("Invalid maxRelativeError: " + c.maxRelError);
+
+        if (c.net.getNumInputArrays() != c.inputs.length)
+            throw new IllegalArgumentException("Invalid input arrays: expect " + c.net.getNumInputArrays() + " inputs");
+        if (c.net.getNumOutputArrays() != c.labels.length)
+            throw new IllegalArgumentException(
+                    "Invalid labels arrays: expect " + c.net.getNumOutputArrays() + " outputs");
+
+        DataType dataType = DataTypeUtil.getDtypeFromContext();
+        if (dataType != DataType.DOUBLE) {
+            throw new IllegalStateException("Cannot perform gradient check: Datatype is not set to double precision ("
+                    + "is: " + dataType + "). Double precision must be used for gradient checks. Set "
+                    + "DataTypeUtil.setDTypeForContext(DataType.DOUBLE); before using GradientCheckUtil");
+        }
+
+        DataType netDataType = c.net.getConfiguration().getDataType();
+        if (netDataType != DataType.DOUBLE) {
+            throw new IllegalStateException("Cannot perform gradient check: Network datatype is not set to double precision ("
+                    + "is: " + netDataType + "). Double precision must be used for gradient checks. Create network with .dataType(DataType.DOUBLE) before using GradientCheckUtil");
+        }
+
+        if(netDataType != c.net.params().dataType()) {
+            throw new IllegalStateException("Parameters datatype does not match network configuration datatype ("
+                    + "is: " + c.net.params().dataType() + "). If network datatype is set to DOUBLE, parameters must also be DOUBLE.");
+        }
+
+        //Check configuration
+        int layerCount = 0;
+        for (String vertexName : c.net.getConfiguration().getVertices().keySet()) {
+            GraphVertex gv = c.net.getConfiguration().getVertices().get(vertexName);
+            if (!(gv instanceof LayerVertex))
+                continue;
+            LayerVertex lv = (LayerVertex) gv;
+
+            if (lv.getLayerConf().getLayer() instanceof BaseLayer) {
+                BaseLayer bl = (BaseLayer) lv.getLayerConf().getLayer();
+                IUpdater u = bl.getIUpdater();
+                if (u instanceof Sgd) {
+                    //Must have LR of 1.0
+                    double lr = ((Sgd) u).getLearningRate();
+                    if (lr != 1.0) {
+                        throw new IllegalStateException("When using SGD updater, must also use lr=1.0 for layer "
+                                + layerCount + "; got " + u + " with lr=" + lr + " for layer \""
+                                + lv.getLayerConf().getLayer().getLayerName() + "\"");
+                    }
+                } else if (!(u instanceof NoOp)) {
+                    throw new IllegalStateException(
+                            "Must have Updater.NONE (or SGD + lr=1.0) for layer " + layerCount + "; got " + u);
+                }
+
+
+            }
+
+            if (lv.getLayerConf().getLayer().getIDropout() != null && c.callEachIter == null) {
+                throw new IllegalStateException("When gradient checking dropout, rng seed must be reset each iteration, or no" +
+                        " dropout should be present during gradient checks - got dropout = "
+                        + lv.getLayerConf().getLayer().getIDropout() + " for layer " + layerCount);
+            }
+        }
+
+        //Set softmax clipping to 0 if necessary, to avoid spurious failures due to clipping
+        for(Layer l : c.net.getLayers()) {
+            if(l instanceof IOutputLayer) {
+                configureLossFnClippingIfPresent((IOutputLayer) l);
+            }
+        }
+
+        for (int i = 0; i < c.inputs.length; i++)
+            c.net.setInput(i, c.inputs[i]);
+        for (int i = 0; i < c.labels.length; i++)
+            c.net.setLabel(i, c.labels[i]);
+
+        c.net.setLayerMaskArrays(c.inputMask, c.labelMask);
+
+        if(c.callEachIter != null){
+            c.callEachIter.accept(c.net);
+        }
+
+        c.net.computeGradientAndScore();
+        Pair<Gradient, Double> gradAndScore = c.net.gradientAndScore();
+
+        ComputationGraphUpdater updater = new ComputationGraphUpdater(c.net);
+        updater.update(gradAndScore.getFirst(), 0, 0, c.net.batchSize(), LayerWorkspaceMgr.noWorkspaces());
+
+        INDArray gradientToCheck = gradAndScore.getFirst().gradient().dup(); //need dup: gradients are a *view* of the full gradient array (which will change every time backprop is done)
+        INDArray originalParams = c.net.params().dup(); //need dup: params are a *view* of full parameters
+
+        val nParams = originalParams.length();
+
+        Map<String, INDArray> paramTable = c.net.paramTable();
+        List<String> paramNames = new ArrayList<>(paramTable.keySet());
+        val paramEnds = new long[paramNames.size()];
+        paramEnds[0] = paramTable.get(paramNames.get(0)).length();
+        for (int i = 1; i < paramEnds.length; i++) {
+            paramEnds[i] = paramEnds[i - 1] + paramTable.get(paramNames.get(i)).length();
+        }
+
+        if(c.excludeParams != null && !c.excludeParams.isEmpty()){
+            log.info("NOTE: parameters will be skipped due to config: {}", c.excludeParams);
+        }
+
+        int currParamNameIdx = 0;
+        int totalNFailures = 0;
+        double maxError = 0.0;
+        MultiDataSet mds = new MultiDataSet(c.inputs, c.labels, c.inputMask, c.labelMask);
+        INDArray params = c.net.params(); //Assumption here: params is a view that we can modify in-place
+        for (long i = 0; i < nParams; i++) {
+            //Get param name
+            if (i >= paramEnds[currParamNameIdx]) {
+                currParamNameIdx++;
+            }
+            String paramName = paramNames.get(currParamNameIdx);
+            if(c.excludeParams != null && c.excludeParams.contains(paramName)){
+                //log.info("Skipping parameters for parameter name: {}", paramName);
+                i = paramEnds[currParamNameIdx++];
+                continue;
+            }
+
+            //(w+epsilon): Do forward pass and score
+            double origValue = params.getDouble(i);
+
+            params.putScalar(i, origValue + c.epsilon);
+            if(c.callEachIter != null) {
+                c.callEachIter.accept(c.net);
+            }
+            double scorePlus = c.net.score(mds, true); //training == true for batch norm, etc (scores and gradients need to be calculated on same thing)
+
+            //(w-epsilon): Do forward pass and score
+            params.putScalar(i, origValue - c.epsilon);
+            if(c.callEachIter != null) {
+                c.callEachIter.accept(c.net);
+            }
+            double scoreMinus = c.net.score(mds, true);
+
+            //Reset original param value
+            params.putScalar(i, origValue);
+
+            //Calculate numerical parameter gradient:
+            double scoreDelta = scorePlus - scoreMinus;
+
+            double numericalGradient = scoreDelta / (2 * c.epsilon);
+            if (Double.isNaN(numericalGradient))
+                throw new IllegalStateException("Numerical gradient was NaN for parameter " + i + " of " + nParams);
+
+            double backpropGradient = gradientToCheck.getDouble(i);
+            //http://cs231n.github.io/neural-networks-3/#gradcheck
+            //use mean centered
+            double relError = Math.abs(backpropGradient - numericalGradient)
+                    / (Math.abs(numericalGradient) + Math.abs(backpropGradient));
+            if (backpropGradient == 0.0 && numericalGradient == 0.0)
+                relError = 0.0; //Edge case: i.e., RNNs with time series length of 1.0
+
+            if (relError > maxError)
+                maxError = relError;
+            if (relError > c.maxRelError || Double.isNaN(relError)) {
+                double absError = Math.abs(backpropGradient - numericalGradient);
+                if (absError < c.minAbsoluteError) {
+                    if(c.print == PrintMode.ALL || c.print == PrintMode.ZEROS && absError == 0.0) {
+                        log.info("Param " + i + " (" + paramName + ") passed: grad= " + backpropGradient
+                                + ", numericalGrad= " + numericalGradient + ", relError= " + relError
+                                + "; absolute error = " + absError + " < minAbsoluteError = " + c.minAbsoluteError);
+                    }
+                } else {
+                    log.info("Param " + i + " (" + paramName + ") FAILED: grad= " + backpropGradient
+                            + ", numericalGrad= " + numericalGradient + ", relError= " + relError
+                            + ", scorePlus=" + scorePlus + ", scoreMinus= " + scoreMinus + ", paramValue = " + origValue);
+                    if (c.exitOnFirstError)
+                        return false;
+                    totalNFailures++;
+                }
+            } else if (c.print == PrintMode.ALL) {
+                log.info("Param " + i + " (" + paramName + ") passed: grad= " + backpropGradient + ", numericalGrad= "
+                        + numericalGradient + ", relError= " + relError);
+            }
+        }
+
+        val nPass = nParams - totalNFailures;
+        log.info("GradientCheckUtil.checkGradients(): " + nParams + " params checked, " + nPass + " passed, "
+                + totalNFailures + " failed. Largest relative error = " + maxError);
+
+        return totalNFailures == 0;
+    }
+
+
+
+    /**
+     * Check backprop gradients for a pretrain layer
+     *
+     * NOTE: gradient checking pretrain layers can be difficult...
+     */
+    public static boolean checkGradientsPretrainLayer(Layer layer, double epsilon, double maxRelError,
+                                                      double minAbsoluteError, boolean print, boolean exitOnFirstError, INDArray input, int rngSeed) {
+
+        LayerWorkspaceMgr mgr = LayerWorkspaceMgr.noWorkspaces();
+
+        //Basic sanity checks on input:
+        if (epsilon <= 0.0 || epsilon > 0.1)
+            throw new IllegalArgumentException("Invalid epsilon: expect epsilon in range (0,0.1], usually 1e-4 or so");
+        if (maxRelError <= 0.0 || maxRelError > 0.25)
+            throw new IllegalArgumentException("Invalid maxRelativeError: " + maxRelError);
+
+        DataType dataType = DataTypeUtil.getDtypeFromContext();
+        if (dataType != DataType.DOUBLE) {
+            throw new IllegalStateException("Cannot perform gradient check: Datatype is not set to double precision ("
+                    + "is: " + dataType + "). Double precision must be used for gradient checks. Set "
+                    + "DataTypeUtil.setDTypeForContext(DataType.DOUBLE); before using GradientCheckUtil");
+        }
+
+        //Check network configuration:
+        layer.setInput(input, LayerWorkspaceMgr.noWorkspaces());
+        Nd4j.getRandom().setSeed(rngSeed);
+        layer.computeGradientAndScore(mgr);
+        Pair<Gradient, Double> gradAndScore = layer.gradientAndScore();
+
+        Updater updater = layer.createUpdater();
+        updater.update(layer, gradAndScore.getFirst(), 0, 0, layer.batchSize(), LayerWorkspaceMgr.noWorkspaces());
+
+        INDArray gradientToCheck = gradAndScore.getFirst().gradient().dup(); //need dup: gradients are a *view* of the full gradient array (which will change every time backprop is done)
+        INDArray originalParams = layer.params().dup(); //need dup: params are a *view* of full parameters
+
+        val nParams = originalParams.length();
+
+        Map<String, INDArray> paramTable = layer.paramTable();
+        List<String> paramNames = new ArrayList<>(paramTable.keySet());
+        val paramEnds = new long[paramNames.size()];
+        paramEnds[0] = paramTable.get(paramNames.get(0)).length();
+        for (int i = 1; i < paramEnds.length; i++) {
+            paramEnds[i] = paramEnds[i - 1] + paramTable.get(paramNames.get(i)).length();
+        }
+
+
+        int totalNFailures = 0;
+        double maxError = 0.0;
+        int currParamNameIdx = 0;
+
+        INDArray params = layer.params(); //Assumption here: params is a view that we can modify in-place
+        for (int i = 0; i < nParams; i++) {
+            //Get param name
+            if (i >= paramEnds[currParamNameIdx]) {
+                currParamNameIdx++;
+            }
+            String paramName = paramNames.get(currParamNameIdx);
+
+            //(w+epsilon): Do forward pass and score
+            double origValue = params.getDouble(i);
+            params.putScalar(i, origValue + epsilon);
+
+            //TODO add a 'score' method that doesn't calculate gradients...
+            Nd4j.getRandom().setSeed(rngSeed);
+            layer.computeGradientAndScore(mgr);
+            double scorePlus = layer.score();
+
+            //(w-epsilon): Do forward pass and score
+            params.putScalar(i, origValue - epsilon);
+            Nd4j.getRandom().setSeed(rngSeed);
+            layer.computeGradientAndScore(mgr);
+            double scoreMinus = layer.score();
+
+            //Reset original param value
+            params.putScalar(i, origValue);
+
+            //Calculate numerical parameter gradient:
+            double scoreDelta = scorePlus - scoreMinus;
+
+            double numericalGradient = scoreDelta / (2 * epsilon);
+            if (Double.isNaN(numericalGradient))
+                throw new IllegalStateException("Numerical gradient was NaN for parameter " + i + " of " + nParams);
+
+            double backpropGradient = gradientToCheck.getDouble(i);
+            //http://cs231n.github.io/neural-networks-3/#gradcheck
+            //use mean centered
+            double relError = Math.abs(backpropGradient - numericalGradient)
+                    / (Math.abs(numericalGradient) + Math.abs(backpropGradient));
+            if (backpropGradient == 0.0 && numericalGradient == 0.0)
+                relError = 0.0; //Edge case: i.e., RNNs with time series length of 1.0
+
+            if (relError > maxError)
+                maxError = relError;
+            if (relError > maxRelError || Double.isNaN(relError)) {
+                double absError = Math.abs(backpropGradient - numericalGradient);
+                if (absError < minAbsoluteError) {
+                    log.info("Param " + i + " (" + paramName + ") passed: grad= " + backpropGradient
+                            + ", numericalGrad= " + numericalGradient + ", relError= " + relError
+                            + "; absolute error = " + absError + " < minAbsoluteError = " + minAbsoluteError);
+                } else {
+                    if (print)
+                        log.info("Param " + i + " (" + paramName + ") FAILED: grad= " + backpropGradient
+                                + ", numericalGrad= " + numericalGradient + ", relError= " + relError
+                                + ", scorePlus=" + scorePlus + ", scoreMinus= " + scoreMinus + ", paramValue = " + origValue);
+                    if (exitOnFirstError)
+                        return false;
+                    totalNFailures++;
+                }
+            } else if (print) {
+                log.info("Param " + i + " (" + paramName + ") passed: grad= " + backpropGradient + ", numericalGrad= "
+                        + numericalGradient + ", relError= " + relError);
+            }
+        }
+
+        if (print) {
+            val nPass = nParams - totalNFailures;
+            log.info("GradientCheckUtil.checkGradients(): " + nParams + " params checked, " + nPass + " passed, "
+                    + totalNFailures + " failed. Largest relative error = " + maxError);
+        }
+
+        return totalNFailures == 0;
+    }
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/adapters/ArgmaxAdapter.java

@@ -0,0 +1,54 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.nn.adapters;
+
+import lombok.val;
+import org.nd4j.adapters.OutputAdapter;
+import org.nd4j.common.base.Preconditions;
+import org.nd4j.linalg.api.ndarray.INDArray;
+import org.nd4j.linalg.factory.Nd4j;
+
+public class ArgmaxAdapter implements OutputAdapter<int[]> {
+
+    /**
+     * This method does conversion from INDArrays to int[], where each element will represents position of the highest element in output INDArray
+     * I.e. Array of {0.25, 0.1, 0.5, 0.15} will return int array with length of 1, and value {2}
+     *
+     * @param outputs
+     * @return
+     */
+    @Override
+    public int[] apply(INDArray... outputs) {
+        Preconditions.checkArgument(outputs.length == 1, "Argmax adapter can have only 1 output");
+        val array = outputs[0];
+        Preconditions.checkArgument(array.rank() < 3, "Argmax adapter requires 2D or 1D output");
+        val result = array.rank() == 2 ? new int[(int) array.size(0)] : new int[1];
+
+        if (array.rank() == 2) {
+            val t = Nd4j.argMax(array, 1);
+            for (int e = 0; e < t.length(); e++)
+                result[e] = (int) t.getDouble(e);
+        } else
+            result[0] = (int) Nd4j.argMax(array, Integer.MAX_VALUE).getDouble(0);
+
+        return result;
+    }
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/adapters/Regression2dAdapter.java

@@ -0,0 +1,48 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.nn.adapters;
+
+import lombok.extern.slf4j.Slf4j;
+import lombok.val;
+import org.nd4j.adapters.OutputAdapter;
+import org.nd4j.common.base.Preconditions;
+import org.nd4j.linalg.api.ndarray.INDArray;
+
+@Slf4j
+public class Regression2dAdapter implements OutputAdapter<double[][]> {
+    @Override
+    public double[][] apply(INDArray... outputs) {
+        Preconditions.checkArgument(outputs.length == 1, "Argmax adapter can have only 1 output");
+        val array = outputs[0];
+        Preconditions.checkArgument(array.rank() < 3, "Argmax adapter requires 2D or 1D output");
+
+        if (array.rank() == 2 && !array.isVector()) {
+            return array.toDoubleMatrix();
+        } else {
+            val result = new double[1][(int) array.length()];
+
+            for (int e = 0; e< array.length(); e++)
+                result[0][e] = array.getDouble(e);
+
+            return result;
+        }
+    }
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/adapters/YoloModelAdapter.java

@@ -0,0 +1,63 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.nn.adapters;
+
+import lombok.AllArgsConstructor;
+import lombok.Builder;
+import lombok.NoArgsConstructor;
+import lombok.val;
+import org.deeplearning4j.nn.api.Model;
+import org.deeplearning4j.nn.api.ModelAdapter;
+import org.deeplearning4j.nn.graph.ComputationGraph;
+import org.deeplearning4j.nn.layers.objdetect.DetectedObject;
+import org.deeplearning4j.nn.layers.objdetect.Yolo2OutputLayer;
+import org.nd4j.linalg.api.ndarray.INDArray;
+import org.nd4j.linalg.exception.ND4JIllegalStateException;
+
+import java.util.List;
+
+@Builder
+@AllArgsConstructor
+@NoArgsConstructor
+public class YoloModelAdapter implements ModelAdapter<List<DetectedObject>> {
+    @Builder.Default private int outputLayerIndex = 0;
+    @Builder.Default private int outputIndex = 0;
+    @Builder.Default private double detectionThreshold = 0.5;
+
+    @Override
+    public List<DetectedObject> apply(Model model, INDArray[] inputs, INDArray[] masks, INDArray[] labelsMasks) {
+        if (model instanceof ComputationGraph) {
+            val blindLayer = ((ComputationGraph) model).getOutputLayer(outputLayerIndex);
+            if (blindLayer instanceof Yolo2OutputLayer) {
+                val output = ((ComputationGraph) model).output(false, inputs, masks, labelsMasks);
+                return ((Yolo2OutputLayer) blindLayer).getPredictedObjects(output[outputIndex], detectionThreshold);
+            } else {
+                throw new ND4JIllegalStateException("Output layer with index [" + outputLayerIndex + "] is NOT Yolo2OutputLayer");
+            }
+        } else
+            throw new ND4JIllegalStateException("Yolo2 model must be ComputationGraph");
+    }
+
+    @Override
+    public List<DetectedObject> apply(INDArray... outputs) {
+        throw new UnsupportedOperationException("Please use apply(Model, INDArray[], INDArray[]) signature");
+    }
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/api/Classifier.java

@@ -0,0 +1,109 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.nn.api;
+
+import org.nd4j.linalg.api.ndarray.INDArray;
+import org.nd4j.linalg.dataset.api.DataSet;
+import org.nd4j.linalg.dataset.api.iterator.DataSetIterator;
+
+import java.util.List;
+
+
+public interface Classifier extends Model {
+
+
+
+    /**
+     * Sets the input and labels and returns a score for the prediction
+     * wrt true labels
+     * @param data the data to score
+     * @return the score for the given input,label pairs
+     */
+    double f1Score(DataSet data);
+
+    /**
+     * Returns the f1 score for the given examples.
+     * Think of this to be like a percentage right.
+     * The higher the number the more it got right.
+     * This is on a scale from 0 to 1.
+     * @param examples te the examples to classify (one example in each row)
+     * @param labels the true labels
+     * @return the scores for each ndarray
+     */
+    double f1Score(INDArray examples, INDArray labels);
+
+    /**
+     * Returns the number of possible labels
+     * @return the number of possible labels for this classifier
+     * @deprecated Will be removed in a future release
+     */
+    @Deprecated
+    int numLabels();
+
+    /**
+     * Train the model based on the datasetiterator
+     * @param iter the iterator to train on
+     */
+    void fit(DataSetIterator iter);
+
+    /**
+     * Takes in a list of examples
+     * For each row, returns a label
+     * @param examples the examples to classify (one example in each row)
+     * @return the labels for each example
+     */
+    int[] predict(INDArray examples);
+
+    /**
+     * Takes in a DataSet of examples
+     * For each row, returns a label
+     * @param dataSet the examples to classify
+     * @return the labels for each example
+     */
+    List<String> predict(DataSet dataSet);
+
+
+    /**
+     * Fit the model
+     * @param examples the examples to classify (one example in each row)
+     * @param labels the example labels(a binary outcome matrix)
+     */
+    void fit(INDArray examples, INDArray labels);
+
+    /**
+     * Fit the model
+     * @param data the data to train on
+     */
+    void fit(DataSet data);
+
+
+
+    /**
+     * Fit the model
+     * @param examples the examples to classify (one example in each row)
+     * @param labels the labels for each example (the number of labels must match
+     *               the number of rows in the example
+     */
+    void fit(INDArray examples, int[] labels);
+
+
+
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/api/FwdPassType.java

@@ -0,0 +1,27 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.nn.api;
+
+public enum FwdPassType {
+    STANDARD,
+    RNN_TIMESTEP,
+    RNN_ACTIVATE_WITH_STORED_STATE
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/api/Layer.java

@@ -0,0 +1,226 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.nn.api;
+
+
+import org.deeplearning4j.nn.conf.CacheMode;
+import org.deeplearning4j.nn.gradient.Gradient;
+import org.deeplearning4j.nn.updater.LayerUpdater;
+import org.deeplearning4j.nn.workspace.LayerWorkspaceMgr;
+import org.deeplearning4j.optimize.api.TrainingListener;
+import org.nd4j.linalg.api.ndarray.INDArray;
+import org.nd4j.common.primitives.Pair;
+
+import java.io.Serializable;
+import java.util.Collection;
+
+public interface Layer extends Serializable, Cloneable, Model, Trainable {
+
+    enum Type {
+        FEED_FORWARD, RECURRENT, CONVOLUTIONAL, CONVOLUTIONAL3D,
+        SUBSAMPLING, UPSAMPLING, RECURSIVE, MULTILAYER, NORMALIZATION
+    }
+
+    enum TrainingMode {
+        TRAIN, TEST
+    }
+
+    default org.deeplearning4j.nn.api.Updater createUpdater() {
+        return new LayerUpdater(this);
+    }
+
+    /**
+     * This method sets given CacheMode for current layer
+     *
+     * @param mode
+     */
+    void setCacheMode(CacheMode mode);
+
+    /**
+     * Calculate the regularization component of the score, for the parameters in this layer<br>
+     * For example, the L1, L2 and/or weight decay components of the loss function<br>
+     *
+     * @param backpropOnlyParams If true: calculate regularization score based on backprop params only. If false: calculate
+     *                           based on all params (including pretrain params, if any)
+     * @return the regularization score of
+     */
+    double calcRegularizationScore(boolean backpropOnlyParams);
+
+    /**
+     * Returns the layer type
+     *
+     * @return
+     */
+    Type type();
+
+
+    /**
+     * Calculate the gradient relative to the error in the next layer
+     *
+     * @param epsilon      w^(L+1)*delta^(L+1). Or, equiv: dC/da, i.e., (dC/dz)*(dz/da) = dC/da, where C
+     *                     is cost function a=sigma(z) is activation.
+     * @param workspaceMgr Workspace manager
+     * @return Pair<Gradient   ,   INDArray> where Gradient is gradient for this layer, INDArray is epsilon (activation gradient)
+     * needed by next layer, but before element-wise multiply by sigmaPrime(z). So for standard feed-forward layer, if this layer is
+     * L, then return.getSecond() == dL/dIn = (w^(L)*(delta^(L))^T)^T. Note that the returned array should be placed in the
+     * {@link org.deeplearning4j.nn.workspace.ArrayType#ACTIVATION_GRAD} workspace via the workspace manager
+     */
+    Pair<Gradient, INDArray> backpropGradient(INDArray epsilon, LayerWorkspaceMgr workspaceMgr);
+
+
+    /**
+     * Perform forward pass and return the activations array with the last set input
+     *
+     * @param training     training or test mode
+     * @param workspaceMgr Workspace manager
+     * @return the activation (layer output) of the last specified input. Note that the returned array should be placed
+     * in the {@link org.deeplearning4j.nn.workspace.ArrayType#ACTIVATIONS} workspace via the workspace manager
+     */
+    INDArray activate(boolean training, LayerWorkspaceMgr workspaceMgr);
+
+    /**
+     * Perform forward pass and return the activations array with the specified input
+     *
+     * @param input    the input to use
+     * @param training train or test mode
+     * @param mgr      Workspace manager.
+     * @return Activations array. Note that the returned array should be placed in the
+     * {@link org.deeplearning4j.nn.workspace.ArrayType#ACTIVATIONS} workspace via the workspace manager
+     */
+    INDArray activate(INDArray input, boolean training, LayerWorkspaceMgr mgr);
+
+    /**
+     * Get the iteration listeners for this layer.
+     */
+    Collection<TrainingListener> getListeners();
+
+    /**
+     * Set the {@link TrainingListener}s for this model. If any listeners have previously been set, they will be
+     * replaced by this method
+     */
+    void setListeners(TrainingListener... listeners);
+
+
+    /**
+     * Set the {@link TrainingListener}s for this model. If any listeners have previously been set, they will be
+     * replaced by this method
+     */
+    void setListeners(Collection<TrainingListener> listeners);
+
+    /**
+     * Set the layer index.
+     */
+    void setIndex(int index);
+
+    /**
+     * Get the layer index.
+     */
+    int getIndex();
+
+    /**
+     * @return The current iteration count (number of parameter updates) for the layer/network
+     */
+    int getIterationCount();
+
+    /**
+     * @return The current epoch count (number of training epochs passed) for the layer/network
+     */
+    int getEpochCount();
+
+    /**
+     * Set the current iteration count (number of parameter updates) for the layer/network
+     */
+    void setIterationCount(int iterationCount);
+
+    /**
+     * Set the current epoch count (number of epochs passed ) for the layer/network
+     */
+    void setEpochCount(int epochCount);
+
+    /**
+     * Set the layer input.
+     */
+    void setInput(INDArray input, LayerWorkspaceMgr workspaceMgr);
+
+    /**
+     * Set current/last input mini-batch size.<br>
+     * Used for score and gradient calculations. Mini batch size may be different from
+     * getInput().size(0) due to reshaping operations - for example, when using RNNs with
+     * DenseLayer and OutputLayer. Called automatically during forward pass.
+     */
+    void setInputMiniBatchSize(int size);
+
+    /**
+     * Get current/last input mini-batch size, as set by setInputMiniBatchSize(int)
+     *
+     * @see Layer#setInputMiniBatchSize(int)
+     */
+    int getInputMiniBatchSize();
+
+    /**
+     * Set the mask array. Note: In general, {@link #feedForwardMaskArray(INDArray, MaskState, int)} should be used in
+     * preference to this.
+     *
+     * @param maskArray Mask array to set
+     */
+    void setMaskArray(INDArray maskArray);
+
+
+    INDArray getMaskArray();
+
+    /**
+     * Returns true if the layer can be trained in an unsupervised/pretrain manner (AE, VAE, etc)
+     *
+     * @return true if the layer can be pretrained (using fit(INDArray), false otherwise
+     */
+    boolean isPretrainLayer();
+
+
+    void clearNoiseWeightParams();
+
+    /**
+     * A performance optimization: mark whether the layer is allowed to modify its input array in-place. In many cases,
+     * this is totally safe - in others, the input array will be shared by multiple layers, and hence it's not safe to
+     * modify the input array.
+     * This is usually used by ops such as dropout.
+     * @param allow If true: the input array is safe to modify. If false: the input array should be copied before it
+     *              is modified (i.e., in-place modifications are un-safe)
+     */
+    void allowInputModification(boolean allow);
+
+
+    /**
+     * Feed forward the input mask array, setting in the layer as appropriate. This allows different layers to
+     * handle masks differently - for example, bidirectional RNNs and normal RNNs operate differently with masks (the
+     * former sets activations to 0 outside of the data present region (and keeps the mask active for future layers like
+     * dense layers), whereas normal RNNs don't zero out the activations/errors )instead relying on backpropagated error
+     * arrays to handle the variable length case.<br>
+     * This is also used for example for networks that contain global pooling layers, arbitrary preprocessors, etc.
+     *
+     * @param maskArray        Mask array to set
+     * @param currentMaskState Current state of the mask - see {@link MaskState}
+     * @param minibatchSize    Current minibatch size. Needs to be known as it cannot always be inferred from the activations
+     *                         array due to reshaping (such as a DenseLayer within a recurrent neural network)
+     * @return New mask array after this layer, along with the new mask state.
+     */
+    Pair<INDArray, MaskState> feedForwardMaskArray(INDArray maskArray, MaskState currentMaskState, int minibatchSize);
+
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/api/MaskState.java

@@ -0,0 +1,25 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.nn.api;
+
+public enum MaskState {
+    Active, Passthrough
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/api/Model.java

@@ -0,0 +1,261 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.nn.api;
+
+import org.deeplearning4j.nn.conf.NeuralNetConfiguration;
+import org.deeplearning4j.nn.gradient.Gradient;
+import org.deeplearning4j.nn.workspace.LayerWorkspaceMgr;
+import org.deeplearning4j.optimize.api.ConvexOptimizer;
+import org.deeplearning4j.optimize.api.TrainingListener;
+import org.nd4j.linalg.api.ndarray.INDArray;
+import org.nd4j.common.primitives.Pair;
+
+import java.util.Collection;
+import java.util.Map;
+
+public interface Model {
+
+
+    org.deeplearning4j.nn.api.Updater createUpdater();
+
+    /**
+     * Init the model
+     */
+    void init();
+
+
+    /**
+     * Set the trainingListeners for the ComputationGraph (and all layers in the network)
+     */
+    void setListeners(Collection<TrainingListener> listeners);
+
+
+    /**
+     * Set the trainingListeners for the ComputationGraph (and all layers in the network)
+     */
+    void setListeners(TrainingListener... listeners);
+
+    /**
+     * This method ADDS additional TrainingListener to existing listeners
+     *
+     * @param listener
+     */
+    void addListeners(TrainingListener... listener);
+
+
+    /**
+     * All models have a fit method
+     */
+    @Deprecated
+    void fit();
+
+    /**
+     * Update layer weights and biases with gradient change
+     */
+    void update(Gradient gradient);
+
+    /**
+     * Perform one update  applying the gradient
+     * @param gradient the gradient to apply
+     */
+    void update(INDArray gradient, String paramType);
+
+
+    /**
+     * The score for the model
+     * @return the score for the model
+     */
+    double score();
+
+
+    /**
+     * Update the score
+     */
+    void computeGradientAndScore(LayerWorkspaceMgr workspaceMgr);
+
+    /**
+     * Parameters of the model (if any)
+     * @return the parameters of the model
+     */
+    INDArray params();
+
+    /**
+     * the number of parameters for the model
+     * @return the number of parameters for the model
+     *
+     */
+    long numParams();
+
+
+    /**
+     * the number of parameters for the model
+     * @return the number of parameters for the model
+     *
+     */
+    long numParams(boolean backwards);
+
+    /**
+     * Set the parameters for this model.
+     * This expects a linear ndarray which then be unpacked internally
+     * relative to the expected ordering of the model
+     * @param params the parameters for the model
+     */
+    void setParams(INDArray params);
+
+    /**
+     * Set the initial parameters array as a view of the full (backprop) network parameters
+     * NOTE: this is intended to be used internally in MultiLayerNetwork and ComputationGraph, not by users.
+     * @param params a 1 x nParams row vector that is a view of the larger (MLN/CG) parameters array
+     */
+    void setParamsViewArray(INDArray params);
+
+
+    INDArray getGradientsViewArray();
+
+    /**
+     * Set the gradients array as a view of the full (backprop) network parameters
+     * NOTE: this is intended to be used internally in MultiLayerNetwork and ComputationGraph, not by users.
+     * @param gradients a 1 x nParams row vector that is a view of the larger (MLN/CG) gradients array
+     */
+    void setBackpropGradientsViewArray(INDArray gradients);
+
+    /**
+     * Fit the model to the given data
+     * @param data the data to fit the model to
+     */
+    void fit(INDArray data, LayerWorkspaceMgr workspaceMgr);
+
+
+    /**
+     * Get the gradient. Note that this method will not calculate the gradient, it will rather return the gradient
+     * that has been computed before.
+     * For calculating the gradient, see {@link Model#computeGradientAndScore(LayerWorkspaceMgr)} } .
+     * @return the gradient for this model, as calculated before
+     */
+    Gradient gradient();
+
+    /**
+     * Get the gradient and score
+     * @return the gradient and score
+     */
+    Pair<Gradient, Double> gradientAndScore();
+
+    /**
+     * The current inputs batch size
+     * @return the current inputs batch size
+     */
+    int batchSize();
+
+
+    /**
+     * The configuration for the neural network
+     * @return the configuration for the neural network
+     */
+    NeuralNetConfiguration conf();
+
+    /**
+     * Setter for the configuration
+     * @param conf
+     */
+    void setConf(NeuralNetConfiguration conf);
+
+    /**
+     * The input/feature matrix for the model
+     * @return the input/feature matrix for the model
+     */
+    INDArray input();
+
+    /**
+     * Returns this models optimizer
+     * @return this models optimizer
+     */
+    ConvexOptimizer getOptimizer();
+
+    /**
+     * Get the parameter
+     * @param param the key of the parameter
+     * @return the parameter vector/matrix with that particular key
+     */
+    INDArray getParam(String param);
+
+    /**
+     * The param table
+     * @return
+     */
+    Map<String, INDArray> paramTable();
+
+    /**
+     * Table of parameters by key, for backprop
+     * For many models (dense layers, etc) - all parameters are backprop parameters
+     * @param backpropParamsOnly If true, return backprop params only. If false: return all params (equivalent to
+     *                           paramsTable())
+     */
+    Map<String, INDArray> paramTable(boolean backpropParamsOnly);
+
+    /**
+     * Setter for the param table
+     * @param paramTable
+     */
+    void setParamTable(Map<String, INDArray> paramTable);
+
+
+    /**
+     * Set the parameter with a new ndarray
+     * @param key the key to se t
+     * @param val the new ndarray
+     */
+    void setParam(String key, INDArray val);
+
+    /**
+     * Clear input
+     */
+    void clear();
+
+
+    /**
+     * Apply any constraints to the model
+     */
+    void applyConstraints(int iteration, int epoch);
+
+
+    void close();
+
+    default void setInput(int inputIndex, INDArray indArray) {
+        throw new UnsupportedOperationException();
+    }
+
+    default void computeGradientAndScore() {
+        throw new UnsupportedOperationException();
+    }
+
+
+
+    //note we do this mostly because layers won't need this most of the time.
+    default void setLabels(int index, INDArray indArray)  {
+        throw new UnsupportedOperationException();
+    }
+
+    default INDArray[] output(INDArray[] input) {
+        throw new UnsupportedOperationException();
+    }
+
+
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/api/ModelAdapter.java

@@ -0,0 +1,32 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.nn.api;
+
+import org.nd4j.adapters.OutputAdapter;
+import org.nd4j.linalg.api.ndarray.INDArray;
+
+public interface ModelAdapter<T> extends OutputAdapter<T> {
+    /**
+     * This method invokes model internally, and does convertion to T
+     * @return
+     */
+    T apply(Model model, INDArray[] inputs, INDArray[] inputMasks, INDArray[] labelsMasks);
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/api/NeuralNetwork.java

@@ -0,0 +1,104 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.nn.api;
+
+import org.deeplearning4j.optimize.api.ConvexOptimizer;
+import org.nd4j.evaluation.IEvaluation;
+import org.nd4j.linalg.api.ndarray.INDArray;
+import org.nd4j.linalg.dataset.api.DataSet;
+import org.nd4j.linalg.dataset.api.MultiDataSet;
+import org.nd4j.linalg.dataset.api.iterator.DataSetIterator;
+import org.nd4j.linalg.dataset.api.iterator.MultiDataSetIterator;
+
+/**
+ * @author raver119
+ */
+public interface NeuralNetwork {
+
+    /**
+     * This method does initialization of model
+     *
+     * PLEASE NOTE: All implementations should track own state, to avoid double spending
+     */
+    void init();
+
+    /**
+     * This method returns model parameters as single INDArray
+     *
+     * @return
+     */
+    INDArray params();
+
+    /**
+     * This method returns updater state (if applicable), null otherwise
+     * @return
+     */
+    INDArray updaterState();
+
+    /**
+     * This method returns Optimizer used for training
+     *
+     * @return
+     */
+    ConvexOptimizer getOptimizer();
+
+    /**
+     * This method fits model with a given DataSet
+     *
+     * @param dataSet
+     */
+    void fit(DataSet dataSet);
+
+    /**
+     * This method fits model with a given MultiDataSet
+     *
+     * @param dataSet
+     */
+    void fit(MultiDataSet dataSet);
+
+    /**
+     * This method fits model with a given DataSetIterator
+     *
+     * @param iterator
+     */
+    void fit(DataSetIterator iterator);
+
+    /**
+     * This method fits model with a given MultiDataSetIterator
+     *
+     * @param iterator
+     */
+    void fit(MultiDataSetIterator iterator);
+
+    /**
+     * This method executes evaluation of the model against given iterator and evaluation implementations
+     *
+     * @param iterator
+     */
+    <T extends IEvaluation> T[] doEvaluation(DataSetIterator iterator, T... evaluations);
+
+    /**
+     * This method executes evaluation of the model against given iterator and evaluation implementations
+     *
+     * @param iterator
+     */
+    <T extends IEvaluation> T[] doEvaluation(MultiDataSetIterator iterator, T... evaluations);
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/api/OptimizationAlgorithm.java

@@ -0,0 +1,30 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.nn.api;
+
+/**
+ * Optimization algorithm to use
+ * @author Adam Gibson
+ *
+ */
+public enum OptimizationAlgorithm {
+    STOCHASTIC_GRADIENT_DESCENT
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/api/ParamInitializer.java

@@ -0,0 +1,105 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.nn.api;
+
+import org.deeplearning4j.nn.conf.NeuralNetConfiguration;
+import org.deeplearning4j.nn.conf.layers.Layer;
+import org.nd4j.linalg.api.ndarray.INDArray;
+
+import java.util.List;
+import java.util.Map;
+
+/**
+ * Param initializer for a layer
+ *
+ * @author Adam Gibson
+ */
+public interface ParamInitializer {
+
+    long numParams(NeuralNetConfiguration conf);
+
+    long numParams(Layer layer);
+
+    /**
+     * Get a list of all parameter keys given the layer configuration
+     *
+     * @param layer Layer
+     * @return All parameter keys
+     */
+    List<String> paramKeys(Layer layer);
+
+    /**
+     * Weight parameter keys given the layer configuration
+     *
+     * @param layer Layer
+     * @return Weight parameter keys
+     */
+    List<String> weightKeys(Layer layer);
+
+    /**
+     * Bias parameter keys given the layer configuration
+     *
+     * @param layer Layer
+     * @return Bias parameter keys
+     */
+    List<String> biasKeys(Layer layer);
+
+    /**
+     * Is the specified parameter a weight?
+     *
+     * @param layer Layer
+     * @param key Key to check
+     * @return True if parameter is a weight
+     */
+    boolean isWeightParam(Layer layer, String key);
+
+    /**
+     * Is the specified parameter a bias?
+     *
+     * @param layer Layer
+     * @param key Key to check
+     * @return True if parameter is a bias
+     */
+    boolean isBiasParam(Layer layer, String key);
+
+    /**
+     * Initialize the parameters
+     *
+     * @param conf             the configuration
+     * @param paramsView       a view of the full network (backprop) parameters
+     * @param initializeParams if true: initialize the parameters according to the configuration. If false: don't modify the
+     *                         values in the paramsView array (but do select out the appropriate subset, reshape etc as required)
+     * @return Map of parameters keyed by type (view of the 'paramsView' array)
+     */
+    Map<String, INDArray> init(NeuralNetConfiguration conf, INDArray paramsView, boolean initializeParams);
+
+    /**
+     * Return a map of gradients (in their standard non-flattened representation), taken from the flattened (row vector) gradientView array.
+     * The idea is that operates in exactly the same way as the paramsView does in {@link #init(Map, NeuralNetConfiguration, INDArray)};
+     * thus the position in the view (and, the array orders) must match those of the parameters
+     *
+     * @param conf         Configuration
+     * @param gradientView The flattened gradients array, as a view of the larger array
+     * @return A map containing an array by parameter type, that is a view of the full network gradients array
+     */
+    Map<String, INDArray> getGradientsFromFlattened(NeuralNetConfiguration conf, INDArray gradientView);
+
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/api/Trainable.java

@@ -0,0 +1,68 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.nn.api;
+
+import org.nd4j.linalg.api.ndarray.INDArray;
+
+import java.util.Map;
+
+public interface Trainable {
+
+    /**
+     * @return Training configuration
+     */
+    TrainingConfig getConfig();
+
+    /**
+     * @return Number of parameters
+     */
+    long numParams();
+
+    /**
+     * @return 1d parameter vector
+     */
+    INDArray params();
+
+    /**
+     * @param backpropOnly If true: return only parameters that are not exclusively used for layerwise pretraining
+     * @return Parameter table
+     */
+    Map<String,INDArray> paramTable(boolean backpropOnly);
+
+    /**
+     * DL4J layers typically produce the sum of the gradients during the backward pass for each layer, and if required
+     * (if minibatch=true) then divide by the minibatch size.<br>
+     * However, there are some exceptions, such as the batch norm mean/variance estimate parameters: these "gradients"
+     * are actually not gradients, but are updates to be applied directly to the parameter vector. Put another way,
+     * most gradients should be divided by the minibatch to get the average; some "gradients" are actually final updates
+     * already, and should not be divided by the minibatch size.
+     *
+     * @param paramName Name of the parameter
+     * @return True if gradients should be divided by minibatch (most params); false otherwise (edge cases like batch norm mean/variance estimates)
+     */
+    boolean updaterDivideByMinibatch(String paramName);
+
+    /**
+     * @return 1D gradients view array
+     */
+    INDArray getGradientsViewArray();
+
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/api/TrainingConfig.java

@@ -0,0 +1,78 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.nn.api;
+
+import org.deeplearning4j.nn.conf.GradientNormalization;
+import org.nd4j.linalg.api.buffer.DataType;
+import org.nd4j.linalg.learning.config.IUpdater;
+import org.nd4j.linalg.learning.regularization.Regularization;
+
+import java.util.List;
+
+public interface TrainingConfig {
+
+    /**
+     * @return Name of the layer
+     */
+    String getLayerName();
+
+    /**
+     * Get the regularization types (l1/l2/weight decay) for the given parameter. Different parameters may have different
+     * regularization types.
+     *
+     * @param paramName Parameter name ("W", "b" etc)
+     * @return Regularization types (if any) for the specified parameter
+     */
+    List<Regularization> getRegularizationByParam(String paramName);
+
+    /**
+     * Is the specified parameter a layerwise pretraining only parameter?<br>
+     * For example, visible bias params in an autoencoder (or, decoder params in a variational autoencoder) aren't
+     * used during supervised backprop.<br>
+     * Layers (like DenseLayer, etc) with no pretrainable parameters will return false for all (valid) inputs.
+     *
+     * @param paramName Parameter name/key
+     * @return True if the parameter is for layerwise pretraining only, false otherwise
+     */
+    boolean isPretrainParam(String paramName);
+
+    /**
+     * Get the updater for the given parameter. Typically the same updater will be used for all updaters, but this
+     * is not necessarily the case
+     *
+     * @param paramName Parameter name
+     * @return IUpdater for the parameter
+     */
+    IUpdater getUpdaterByParam(String paramName);
+
+    /**
+     * @return The gradient normalization configuration
+     */
+    GradientNormalization getGradientNormalization();
+
+    /**
+     * @return The gradient normalization threshold
+     */
+    double getGradientNormalizationThreshold();
+
+    void setDataType(DataType dataType);
+
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/api/Updater.java

@@ -0,0 +1,58 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.nn.api;
+
+import org.deeplearning4j.nn.gradient.Gradient;
+import org.nd4j.linalg.api.ndarray.INDArray;
+import org.deeplearning4j.nn.workspace.LayerWorkspaceMgr;
+
+import java.io.Serializable;
+
+/**
+ * Update the model
+ *
+ * @author Adam Gibson
+ */
+public interface Updater extends Serializable {
+
+    /**
+     * Set the internal (historical) state view array for this updater
+     *
+     * @param layer      Layer that this updater belongs to
+     * @param viewArray  View array
+     * @param initialize Whether to initialize the array or not
+     */
+    void setStateViewArray(Trainable layer, INDArray viewArray, boolean initialize);
+
+    /**
+     * @return the view array for this updater
+     */
+    INDArray getStateViewArray();
+
+    /**
+     * Updater: updates the model
+     *
+     * @param layer
+     * @param gradient
+     * @param iteration
+     */
+    void update(Trainable layer, Gradient gradient, int iteration, int epoch, int miniBatchSize, LayerWorkspaceMgr workspaceMgr);
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/api/layers/IOutputLayer.java

@@ -0,0 +1,70 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.nn.api.layers;
+
+import org.deeplearning4j.nn.api.Classifier;
+import org.deeplearning4j.nn.api.Layer;
+import org.nd4j.linalg.api.ndarray.INDArray;
+import org.deeplearning4j.nn.workspace.LayerWorkspaceMgr;
+
+public interface IOutputLayer extends Layer, Classifier {
+
+    /**
+     * Returns true if labels are required
+     * for this output layer
+     * @return true if this output layer needs labels or not
+     */
+    boolean needsLabels();
+
+    /**
+     * Set the labels array for this output layer
+     *
+     * @param labels Labels array to set
+     */
+    void setLabels(INDArray labels);
+
+    /**
+     * Get the labels array previously set with {@link #setLabels(INDArray)}
+     *
+     * @return Labels array, or null if it has not been set
+     */
+    INDArray getLabels();
+
+    /**
+     * Compute score after labels and input have been set.
+     *
+     * @param fullNetworkRegScore Regularization score (l1/l2/weight decay) for the entire network
+     * @param training            whether score should be calculated at train or test time (this affects things like application of
+     *                            dropout, etc)
+     * @return score (loss function)
+     */
+    double computeScore(double fullNetworkRegScore, boolean training, LayerWorkspaceMgr workspaceMgr);
+
+    /**
+     * Compute the score for each example individually, after labels and input have been set.
+     *
+     * @param fullNetworkRegScore Regularization score (l1/l2/weight decay) for the entire network
+     * @return A column INDArray of shape [numExamples,1], where entry i is the score of the ith example
+     */
+    INDArray computeScoreForExamples(double fullNetworkRegScore, LayerWorkspaceMgr workspaceMgr);
+
+
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/api/layers/LayerConstraint.java

@@ -0,0 +1,56 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.nn.api.layers;
+
+import org.deeplearning4j.nn.api.Layer;
+import org.nd4j.shade.jackson.annotation.JsonTypeInfo;
+
+import java.io.Serializable;
+import java.util.Set;
+
+@JsonTypeInfo(use = JsonTypeInfo.Id.CLASS, include = JsonTypeInfo.As.PROPERTY, property = "@class")
+public interface LayerConstraint extends Cloneable, Serializable {
+
+    /**
+     * Apply a given constraint to a layer at each iteration
+     * in the provided epoch, after parameters have been updated.
+     *
+     * @param layer org.deeplearning4j.nn.api.Layer
+     * @param iteration given iteration as integer
+     * @param epoch current epoch as integer
+     */
+    void applyConstraint(Layer layer, int iteration, int epoch);
+
+    /**
+     * Set the parameters that this layer constraint should be applied to
+     *
+     * @param params Parameters that the layer constraint should be applied to
+     */
+    void setParams(Set<String> params);
+
+    /**
+     * @return Set of parameters that this layer constraint will be applied to
+     */
+    Set<String> getParams();
+
+    LayerConstraint clone();
+
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/api/layers/RecurrentLayer.java

@@ -0,0 +1,103 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.nn.api.layers;
+
+import org.deeplearning4j.nn.api.Layer;
+import org.deeplearning4j.nn.gradient.Gradient;
+import org.nd4j.linalg.api.ndarray.INDArray;
+import org.nd4j.common.primitives.Pair;
+import org.deeplearning4j.nn.workspace.LayerWorkspaceMgr;
+
+import java.util.Map;
+
+public interface RecurrentLayer extends Layer {
+
+    /**
+     * Do one or more time steps using the previous time step state stored in stateMap.<br>
+     * Can be used to efficiently do forward pass one or n-steps at a time (instead of doing
+     * forward pass always from t=0)<br>
+     * If stateMap is empty, default initialization (usually zeros) is used<br>
+     * Implementations also update stateMap at the end of this method
+     *
+     * @param input Input to this layer
+     * @return activations
+     */
+    INDArray rnnTimeStep(INDArray input, LayerWorkspaceMgr workspaceMgr);
+
+    /**
+     * Returns a shallow copy of the RNN stateMap (that contains the stored history for use in methods such
+     * as rnnTimeStep
+     */
+    Map<String, INDArray> rnnGetPreviousState();
+
+    /**
+     * Set the stateMap (stored history). Values set using this method will be used in next call to rnnTimeStep()
+     */
+    void rnnSetPreviousState(Map<String, INDArray> stateMap);
+
+    /**
+     * Reset/clear the stateMap for rnnTimeStep() and tBpttStateMap for rnnActivateUsingStoredState()
+     */
+    void rnnClearPreviousState();
+
+    /**
+     * Similar to rnnTimeStep, this method is used for activations using the state
+     * stored in the stateMap as the initialization. However, unlike rnnTimeStep this
+     * method does not alter the stateMap; therefore, unlike rnnTimeStep, multiple calls to
+     * this method (with identical input) will:<br>
+     * (a) result in the same output<br>
+     * (b) leave the state maps (both stateMap and tBpttStateMap) in an identical state
+     *
+     * @param input             Layer input
+     * @param training          if true: training. Otherwise: test
+     * @param storeLastForTBPTT If true: store the final state in tBpttStateMap for use in truncated BPTT training
+     * @return Layer activations
+     */
+    INDArray rnnActivateUsingStoredState(INDArray input, boolean training, boolean storeLastForTBPTT, LayerWorkspaceMgr workspaceMg);
+
+    /**
+     * Get the RNN truncated backpropagations through time (TBPTT) state for the recurrent layer.
+     * The TBPTT state is used to store intermediate activations/state between updating parameters when doing
+     * TBPTT learning
+     *
+     * @return State for the RNN layer
+     */
+    Map<String, INDArray> rnnGetTBPTTState();
+
+    /**
+     * Set the RNN truncated backpropagations through time (TBPTT) state for the recurrent layer.
+     * The TBPTT state is used to store intermediate activations/state between updating parameters when doing
+     * TBPTT learning
+     *
+     * @param state TBPTT state to set
+     */
+    void rnnSetTBPTTState(Map<String, INDArray> state);
+
+    /**
+     * Truncated BPTT equivalent of Layer.backpropGradient().
+     * Primary difference here is that forward pass in the context of BPTT is that we do
+     * forward pass using stored state for truncated BPTT vs. from zero initialization
+     * for standard BPTT.
+     */
+    Pair<Gradient, INDArray> tbpttBackpropGradient(INDArray epsilon, int tbpttBackLength, LayerWorkspaceMgr workspaceMgr);
+
+
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/conf/BackpropType.java

@@ -0,0 +1,30 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.nn.conf;
+
+public enum BackpropType {
+    /** Default option. Used for training most networks, including MLP, DBNs, CNNs etc.*/
+    Standard,
+    /** Truncated BackPropagation Through Time. Only applicable in context of
+     * training networks with recurrent neural network layers such as GravesLSTM
+     */
+    TruncatedBPTT
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/conf/BaseBuilder.java

@@ -0,0 +1,219 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.nn.conf;
+
+import lombok.Data;
+import lombok.NonNull;
+import org.deeplearning4j.nn.conf.inputs.InputType;
+import org.nd4j.linalg.api.buffer.DataType;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+@Data
+public abstract class BaseBuilder {
+
+    protected static final int DEFAULT_TBPTT_LENGTH = 20;
+
+    protected List<NeuralNetConfiguration> confs = new ArrayList<>();
+    protected double dampingFactor = 100;
+    protected Map<Integer, InputPreProcessor> inputPreProcessors = new HashMap<>();
+    protected BackpropType backpropType = BackpropType.Standard;
+    protected int tbpttFwdLength = DEFAULT_TBPTT_LENGTH;
+    protected int tbpttBackLength = DEFAULT_TBPTT_LENGTH;
+    protected InputType inputType;
+
+    protected WorkspaceMode trainingWorkspaceMode = WorkspaceMode.ENABLED;
+    protected WorkspaceMode inferenceWorkspaceMode = WorkspaceMode.ENABLED;
+    protected CacheMode cacheMode = CacheMode.NONE;
+    protected boolean validateOutputConfig = true;
+    protected boolean validateTbpttConfig = true;
+    protected DataType dataType;
+    protected boolean overrideNinUponBuild = true;
+
+
+    /**
+     * Whether to over ride the nIn
+     * configuration forcibly upon construction.
+     * Default value is true
+     * @param overrideNinUponBuild Whether to over ride the nIn
+     *           configuration forcibly upon construction.
+     * @return builder pattern
+     */
+    public  <T extends BaseBuilder> T overrideNinUponBuild(boolean overrideNinUponBuild) {
+        this.overrideNinUponBuild = overrideNinUponBuild;
+        return  (T) this;
+    }
+
+    /**
+     * Specify the processors.
+     * These are used at each layer for doing things like normalization and
+     * shaping of input.
+     *
+     * @param processor what to use to preProcess the data.
+     * @return builder pattern
+     */
+    public  <T extends BaseBuilder> T inputPreProcessor(Integer layer, InputPreProcessor processor) {
+        inputPreProcessors.put(layer, processor);
+        return  (T) this;
+    }
+
+    public  <T extends BaseBuilder> T inputPreProcessors(Map<Integer, InputPreProcessor> processors) {
+        this.inputPreProcessors = processors;
+        return  (T) this;
+    }
+
+    /**
+     * @deprecated Use {@link NeuralNetConfiguration.Builder#trainingWorkspaceMode(WorkspaceMode)}
+     */
+    @Deprecated
+    public  <T extends BaseBuilder> T trainingWorkspaceMode(@NonNull WorkspaceMode workspaceMode) {
+        this.trainingWorkspaceMode = workspaceMode;
+        return  (T) this;
+    }
+
+    /**
+     * @deprecated Use {@link NeuralNetConfiguration.Builder#inferenceWorkspaceMode(WorkspaceMode)}
+     */
+    @Deprecated
+    public  <T extends BaseBuilder> T inferenceWorkspaceMode(@NonNull WorkspaceMode workspaceMode) {
+        this.inferenceWorkspaceMode = workspaceMode;
+        return (T) this;
+    }
+
+    /**
+     * This method defines how/if preOutput cache is handled:
+     * NONE: cache disabled (default value)
+     * HOST: Host memory will be used
+     * DEVICE: GPU memory will be used (on CPU backends effect will be the same as for HOST)
+     *
+     * @param cacheMode
+     * @return
+     */
+    public  <T extends BaseBuilder> T cacheMode(@NonNull CacheMode cacheMode) {
+        this.cacheMode = cacheMode;
+        return  (T) this;
+    }
+
+    /**
+     * The type of backprop. Default setting is used for most networks (MLP, CNN etc),
+     * but optionally truncated BPTT can be used for training recurrent neural networks.
+     * If using TruncatedBPTT make sure you set both tBPTTForwardLength() and tBPTTBackwardLength()
+     */
+    public  <T extends BaseBuilder> T backpropType(@NonNull BackpropType type) {
+        this.backpropType = type;
+        return  (T) this;
+    }
+
+    /**
+     * When doing truncated BPTT: how many steps should we do?<br>
+     * Only applicable when doing backpropType(BackpropType.TruncatedBPTT)<br>
+     * See: <a href="http://www.cs.utoronto.ca/~ilya/pubs/ilya_sutskever_phd_thesis.pdf">http://www.cs.utoronto.ca/~ilya/pubs/ilya_sutskever_phd_thesis.pdf</a>
+     *
+     * @param bpttLength length > 0
+     */
+    public  <T extends BaseBuilder> T tBPTTLength(int bpttLength) {
+        tBPTTForwardLength(bpttLength);
+        return tBPTTBackwardLength(bpttLength);
+    }
+
+    /**
+     * When doing truncated BPTT: how many steps of forward pass should we do
+     * before doing (truncated) backprop?<br>
+     * Only applicable when doing backpropType(BackpropType.TruncatedBPTT)<br>
+     * Typically tBPTTForwardLength parameter is same as the tBPTTBackwardLength parameter,
+     * but may be larger than it in some circumstances (but never smaller)<br>
+     * Ideally your training data time series length should be divisible by this
+     * This is the k1 parameter on pg23 of
+     * <a href="http://www.cs.utoronto.ca/~ilya/pubs/ilya_sutskever_phd_thesis.pdf">http://www.cs.utoronto.ca/~ilya/pubs/ilya_sutskever_phd_thesis.pdf</a>
+     *
+     * @param forwardLength Forward length > 0, >= backwardLength
+     */
+    public  <T extends BaseBuilder> T tBPTTForwardLength(int forwardLength) {
+        this.tbpttFwdLength = forwardLength;
+        return (T) this;
+    }
+
+    /**
+     * When doing truncated BPTT: how many steps of backward should we do?<br>
+     * Only applicable when doing backpropType(BackpropType.TruncatedBPTT)<br>
+     * This is the k2 parameter on pg23 of
+     * <a href="http://www.cs.utoronto.ca/~ilya/pubs/ilya_sutskever_phd_thesis.pdf">http://www.cs.utoronto.ca/~ilya/pubs/ilya_sutskever_phd_thesis.pdf</a>
+     *
+     * @param backwardLength <= forwardLength
+     */
+    public  <T extends BaseBuilder> T tBPTTBackwardLength(int backwardLength) {
+        this.tbpttBackLength = backwardLength;
+        return (T) this;
+    }
+
+    public  <T extends BaseBuilder> T confs(List<NeuralNetConfiguration> confs) {
+        this.confs = confs;
+        return (T) this;
+    }
+
+    public  <T extends BaseBuilder> T setInputType(InputType inputType) {
+        this.inputType = inputType;
+        return (T) this;
+    }
+
+    /**
+     * Enabled by default. If enabled, the output layer configuration will be validated, to throw an exception on
+     * likely invalid outputs - such as softmax + nOut=1, or LossMCXENT + Tanh.<br>
+     * If disabled (false) no output layer validation will be performed.<br>
+     * Disabling this validation is not recommended, as the configurations that fail validation usually will
+     * not be able to learn correctly. However, the option to disable this validation is provided for advanced users
+     * when creating non-standard architectures.
+     *
+     * @param validate If true: validate output layer configuration. False: don't validate
+     */
+    public <T extends BaseBuilder> T validateOutputLayerConfig(boolean validate) {
+        this.validateOutputConfig = validate;
+        return (T) this;
+    }
+
+    /**
+     * Enabled by default. If enabled, an exception will be throw when using the (invalid) combination of truncated
+     * backpropagation through time (TBPTT) with either a GlobalPoolingLayer or LastTimeStepLayer.<br>
+     * It is possible to disable this validation to allow what is almost certainly an invalid configuration to be used,
+     * however this is not recommended.
+     *
+     * @param validate Whether TBPTT validation should be performed
+     */
+    public <T extends BaseBuilder> T  validateTbpttConfig(boolean validate){
+        this.validateTbpttConfig = validate;
+        return (T) this;
+    }
+
+    /**
+     * Set the DataType for the network parameters and activations for all layers in the network. Default: Float
+     * @param dataType Datatype to use for parameters and activations
+     */
+    public <T extends BaseBuilder> T  dataType(@NonNull DataType dataType) {
+        this.dataType = dataType;
+        return (T) this;
+    }
+
+    public abstract <T> T build();
+
+
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/conf/CNN2DFormat.java

@@ -0,0 +1,42 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.nn.conf;
+
+public enum CNN2DFormat implements DataFormat {
+    NCHW,
+    NHWC;
+
+    /**
+     * Returns a string that explains the dimensions:<br>
+     * NCHW -> returns "[minibatch, channels, height, width]"<br>
+     * NHWC -> returns "[minibatch, height, width, channels]"
+     */
+    public String dimensionNames(){
+        switch (this){
+            case NCHW:
+                return "[minibatch, channels, height, width]";
+            case NHWC:
+                return "[minibatch, height, width, channels]";
+            default:
+                throw new IllegalStateException("Unknown enum: " + this);   //Should never happen
+        }
+    }
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/conf/CacheMode.java

@@ -0,0 +1,38 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.nn.conf;
+
+public enum CacheMode {
+    /**
+     * Device memory will be used for cache (if current backend support such differentiation)
+     */
+    DEVICE,
+
+    /**
+     * Host memory will be used for cache
+     */
+    HOST,
+
+    /**
+     * Cache won't be used during training
+     */
+    NONE
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/conf/ConfClassLoading.java

@@ -0,0 +1,269 @@
+
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.nn.conf;
+
+import org.deeplearning4j.nn.conf.constraint.MaxNormConstraint;
+import org.deeplearning4j.nn.conf.constraint.MinMaxNormConstraint;
+import org.deeplearning4j.nn.conf.constraint.NonNegativeConstraint;
+import org.deeplearning4j.nn.conf.constraint.UnitNormConstraint;
+import org.deeplearning4j.nn.conf.distribution.*;
+import org.deeplearning4j.nn.conf.dropout.AlphaDropout;
+import org.deeplearning4j.nn.conf.dropout.GaussianDropout;
+import org.deeplearning4j.nn.conf.dropout.GaussianNoise;
+import org.deeplearning4j.nn.conf.dropout.SpatialDropout;
+import org.deeplearning4j.nn.conf.graph.*;
+import org.deeplearning4j.nn.conf.graph.rnn.DuplicateToTimeSeriesVertex;
+import org.deeplearning4j.nn.conf.graph.rnn.LastTimeStepVertex;
+import org.deeplearning4j.nn.conf.layers.*;
+import org.deeplearning4j.nn.conf.layers.CnnLossLayer;
+import org.deeplearning4j.nn.conf.layers.Convolution1DLayer;
+import org.deeplearning4j.nn.conf.layers.ConvolutionLayer;
+import org.deeplearning4j.nn.conf.layers.ZeroPadding1DLayer;
+import org.deeplearning4j.nn.conf.layers.ZeroPadding3DLayer;
+import org.deeplearning4j.nn.conf.layers.ZeroPaddingLayer;
+import org.deeplearning4j.nn.conf.layers.convolutional.Cropping1D;
+import org.deeplearning4j.nn.conf.layers.convolutional.Cropping2D;
+import org.deeplearning4j.nn.conf.layers.convolutional.Cropping3D;
+import org.deeplearning4j.nn.conf.layers.misc.ElementWiseMultiplicationLayer;
+import org.deeplearning4j.nn.conf.layers.recurrent.Bidirectional;
+import org.deeplearning4j.nn.conf.layers.recurrent.LastTimeStep;
+import org.deeplearning4j.nn.conf.layers.recurrent.SimpleRnn;
+import org.deeplearning4j.nn.conf.layers.recurrent.TimeDistributed;
+import org.deeplearning4j.nn.conf.layers.samediff.SameDiffLambdaLayer;
+import org.deeplearning4j.nn.conf.layers.samediff.SameDiffLambdaVertex;
+import org.deeplearning4j.nn.conf.layers.samediff.SameDiffLayer;
+import org.deeplearning4j.nn.conf.layers.samediff.SameDiffOutputLayer;
+import org.deeplearning4j.nn.conf.layers.util.MaskZeroLayer;
+import org.deeplearning4j.nn.conf.layers.variational.VariationalAutoencoder;
+import org.deeplearning4j.nn.conf.ocnn.OCNNOutputLayer;
+import org.deeplearning4j.nn.graph.ComputationGraph;
+import org.deeplearning4j.nn.layers.FrozenLayer;
+import org.deeplearning4j.nn.layers.RepeatVector;
+import org.deeplearning4j.nn.layers.convolution.*;
+import org.deeplearning4j.nn.layers.objdetect.Yolo2OutputLayer;
+import org.deeplearning4j.nn.layers.util.IdentityLayer;
+import org.deeplearning4j.nn.layers.util.MaskLayer;
+import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;
+import org.nd4j.autodiff.functions.DifferentialFunction;
+import org.nd4j.common.primitives.AtomicBoolean;
+import org.nd4j.common.tools.ClassInitializerUtil;
+import org.nd4j.linalg.activations.impl.*;
+import org.nd4j.linalg.api.ops.impl.layers.convolution.DepthToSpace;
+import org.nd4j.linalg.api.ops.impl.transforms.custom.BatchToSpace;
+import org.nd4j.linalg.lossfunctions.ILossFunction;
+import org.nd4j.linalg.lossfunctions.LossFunctions;
+import org.nd4j.linalg.lossfunctions.impl.*;
+
+public class ConfClassLoading {
+    private static AtomicBoolean invoked = new AtomicBoolean(false);
+
+    public static void loadConfigClasses() throws ClassNotFoundException {
+        if(invoked.get()) return;
+
+        ClassInitializerUtil.tryLoadClasses(MultiLayerConfiguration.class,
+                MultiLayerConfiguration.Builder.class,
+                LossFunctions.class,
+                ILossFunction.class,
+                LossMSE.class,
+                LossMAE.class,
+                LossBinaryXENT.class,
+                LossFMeasure.class,
+                LossSparseMCXENT.class,
+                LossNegativeLogLikelihood.class,
+                LossMCXENT.class,
+                LossKLD.class,
+                LossL1.class,
+                LossL2.class,
+                LossHinge.class,
+                LossSquaredHinge.class,
+                LossCosineProximity.class,
+                LossPoisson.class,
+                LossMAPE.class,
+                LossMSLE.class,
+                LossL2.class,
+                LossL1.class,
+                LossWasserstein.class,
+                MultiLayerNetwork.class,
+                NeuralNetConfiguration.class,
+                NeuralNetConfiguration.Builder.class,
+                ComputationGraphConfiguration.class,
+                ComputationGraphConfiguration.GraphBuilder.class,
+                ComputationGraph.class,
+                Layer.class,
+                Layer.Builder.class,
+                FeedForwardLayer.class,
+                BaseOutputLayer.class,
+                BaseLayer.class,
+                ConvolutionLayer.class,
+                ConvolutionLayer.Builder.class,
+                Convolution1DLayer.class,
+                Convolution1DLayer.Builder.class,
+                Convolution3DLayer.class,
+                Class.forName("org.deeplearning4j.nn.conf.layers.SubsamplingLayer$1"),
+                org.nd4j.linalg.util.LongUtils.class,
+                DifferentialFunction.class,
+                ConvolutionMode.class,
+                CNN2DFormat.class,
+                PoolingType.class,
+                SubsamplingLayer.class,
+                SubsamplingLayer.Builder.class,
+                PrimaryCapsules.class,
+                CapsuleLayer.class,
+                RecurrentAttentionLayer.class,
+                //activations,
+                ActivationCube.class,
+                ActivationELU.class,
+                ActivationHardSigmoid.class,
+                ActivationHardTanH.class,
+                ActivationIdentity.class,
+                ActivationLReLU.class,
+                ActivationRationalTanh.class,
+                ActivationRectifiedTanh.class,
+                ActivationReLU.class,
+                ActivationReLU6.class,
+                ActivationSELU.class,
+                ActivationSwish.class,
+                ActivationRReLU.class,
+                ActivationSigmoid.class,
+                ActivationSoftmax.class,
+                ActivationSoftPlus.class,
+                ActivationSoftSign.class,
+                ActivationTanH.class,
+                ActivationThresholdedReLU.class,
+                ActivationGELU.class,
+                ActivationMish.class,
+
+
+
+                //normalizations
+                MaxNormConstraint.class,
+                MinMaxNormConstraint.class,
+                NonNegativeConstraint.class,
+                UnitNormConstraint.class,
+                //distributions
+                BinomialDistribution.class,
+                ConstantDistribution.class,
+                LogNormalDistribution.class,
+                NormalDistribution.class,
+                OrthogonalDistribution.class,
+                TruncatedNormalDistribution.class,
+                UniformDistribution.class,
+
+                //vertices:
+                AttentionVertex.class,
+                DotProductAttentionLayer.class,
+                ElementWiseVertex.class,
+                GraphVertex.class,
+                L2Vertex.class,
+                MergeVertex.class,
+                PreprocessorVertex.class,
+                ReshapeVertex.class,
+                ScaleVertex.class,
+                ShiftVertex.class,
+                SubsetVertex.class,
+                UnstackVertex.class,
+                StackVertex.class,
+                LastTimeStepVertex.class,
+                DuplicateToTimeSeriesVertex.class,
+                PreprocessorVertex.class,
+
+                //samediff
+                SameDiffLambdaLayer.class,
+                SameDiffLambdaVertex.class,
+                SameDiffLayer.class,
+                SameDiffOutputLayer.class,
+
+
+
+                //dropout
+                AlphaDropout.class,
+                GaussianDropout.class,
+                GaussianNoise.class,
+                SpatialDropout.class,
+
+                //layers
+                DenseLayer.class,
+                AutoEncoder.class,
+                VariationalAutoencoder.class,
+                ElementWiseMultiplicationLayer.class,
+                PReLULayer.class,
+                EmbeddingLayer.class,
+                OutputLayer.class,
+                EmbeddingSequenceLayer.class,
+                BatchNormalization.class,
+                LocalResponseNormalization.class,
+                Yolo2OutputLayer.class,
+                IdentityLayer.class,
+                MaskLayer.class,
+                OCNNOutputLayer.class,
+                GlobalPoolingLayer.class,
+                LastTimeStep.class,
+                MaskZeroLayer.class,
+                SimpleRnn.class,
+                TimeDistributed.class,
+                Bidirectional.class,
+                ActivationLayer.class,
+                DropoutLayer.class,
+                FrozenLayer.class,
+                RepeatVector.class,
+                Subsampling1DLayer.class,
+                Subsampling3DLayer.class,
+                Convolution1DLayer.class,
+                Convolution3DLayer.class,
+                ConvolutionLayer.class,
+                Upsampling1D.class,
+                Upsampling2D.class,
+                Upsampling3D.class,
+                Deconvolution2D.class,
+                Deconvolution3D.class,
+                CnnLossLayer.class,
+                CenterLossOutputLayer.class,
+                RnnOutputLayer.class,
+                OutputLayer.class,
+                LastTimeStep.class,
+                Cropping1DLayer.class,
+                Cropping2DLayer.class,
+                Cropping3DLayer.class,
+                Cropping1D.class,
+                Cropping2D.class,
+                Cropping3D.class,
+                SeparableConvolution2DLayer.class,
+                ZeroPadding1DLayer.class,
+                ZeroPadding3DLayer.class,
+                ZeroPaddingLayer.class,
+                SpaceToBatch.class,
+                SpaceToDepth.class,
+                BatchToSpace.class,
+                DepthToSpace.class,
+                DepthwiseConvolution2D.class);
+    }
+
+
+    static {
+        try {
+            loadConfigClasses();
+        } catch (ClassNotFoundException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/conf/ConvolutionMode.java

@@ -0,0 +1,45 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.nn.conf;
+
+import org.nd4j.linalg.api.ops.impl.layers.convolution.config.PaddingMode;
+import org.nd4j.linalg.api.ops.impl.transforms.Pad;
+
+public enum ConvolutionMode {
+
+    Strict, Truncate, Same, Causal;
+
+
+    public static PaddingMode mapToMode(ConvolutionMode convolutionMode) {
+        switch(convolutionMode) {
+            case Strict:
+            case Truncate:
+                return PaddingMode.VALID;
+            case Same:
+                return PaddingMode.SAME;
+            case Causal:
+                return PaddingMode.CAUSAL;
+            default:
+                throw new IllegalArgumentException("No convolution mode found!");
+        }
+    }
+
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/conf/DataFormat.java

@@ -0,0 +1,30 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.nn.conf;
+
+import org.deeplearning4j.nn.conf.serde.format.DataFormatDeserializer;
+import org.deeplearning4j.nn.conf.serde.format.DataFormatSerializer;
+import org.nd4j.shade.jackson.databind.annotation.JsonDeserialize;
+import org.nd4j.shade.jackson.databind.annotation.JsonSerialize;
+
+@JsonSerialize(using = DataFormatSerializer.class)
+@JsonDeserialize(using = DataFormatDeserializer.class)
+public interface DataFormat {
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/conf/GradientNormalization.java

@@ -0,0 +1,25 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.nn.conf;
+
+public enum GradientNormalization {
+    None, RenormalizeL2PerLayer, RenormalizeL2PerParamType, ClipElementWiseAbsoluteValue, ClipL2PerLayer, ClipL2PerParamType
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/conf/InputPreProcessor.java

@@ -0,0 +1,69 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.nn.conf;
+
+
+import org.deeplearning4j.nn.api.MaskState;
+import org.deeplearning4j.nn.conf.inputs.InputType;
+import org.nd4j.linalg.api.ndarray.INDArray;
+import org.nd4j.common.primitives.Pair;
+import org.deeplearning4j.nn.workspace.LayerWorkspaceMgr;
+import org.nd4j.shade.jackson.annotation.JsonTypeInfo;
+
+import java.io.Serializable;
+
+@JsonTypeInfo(use = JsonTypeInfo.Id.CLASS, include = JsonTypeInfo.As.PROPERTY, property = "@class")
+public interface InputPreProcessor extends Serializable, Cloneable {
+
+    /**
+     * Pre preProcess input/activations for a multi layer network
+     * @param input the input to pre preProcess
+     * @param miniBatchSize Minibatch size
+     * @param workspaceMgr Workspace manager
+     * @return the processed input. Note that the returned array should be placed in the
+     *         {@link org.deeplearning4j.nn.workspace.ArrayType#ACTIVATIONS} workspace via the workspace manager
+     */
+    INDArray preProcess(INDArray input, int miniBatchSize, LayerWorkspaceMgr workspaceMgr);
+
+    /**Reverse the preProcess during backprop. Process Gradient/epsilons before
+     * passing them to the layer below.
+     * @param output which is a pair of the gradient and epsilon
+     * @param miniBatchSize Minibatch size
+     * @param workspaceMgr Workspace manager
+     * @return the reverse of the pre preProcess step (if any). Note that the returned array should be
+     *         placed in {@link org.deeplearning4j.nn.workspace.ArrayType#ACTIVATION_GRAD} workspace via the
+     *         workspace manager
+     */
+    INDArray backprop(INDArray output, int miniBatchSize, LayerWorkspaceMgr workspaceMgr);
+
+    InputPreProcessor clone();
+
+    /**
+     * For a given type of input to this preprocessor, what is the type of the output?
+     *
+     * @param inputType    Type of input for the preprocessor
+     * @return             Type of input after applying the preprocessor
+     */
+    InputType getOutputType(InputType inputType);
+
+
+    Pair<INDArray, MaskState> feedForwardMaskArray(INDArray maskArray, MaskState currentMaskState, int minibatchSize);
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/conf/ListBuilder.java

@@ -0,0 +1,234 @@
+package org.deeplearning4j.nn.conf;
+
+import lombok.Data;
+import lombok.NonNull;
+import lombok.extern.slf4j.Slf4j;
+import org.deeplearning4j.nn.conf.inputs.InputType;
+import org.deeplearning4j.nn.conf.layers.Layer;
+import org.nd4j.common.base.Preconditions;
+import org.nd4j.linalg.api.buffer.DataType;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+/**
+ * Fluent interface for building a list of configurations
+ */
+@Slf4j
+@Data
+public class ListBuilder extends BaseBuilder {
+    private int layerCounter = -1; //Used only for .layer(Layer) method
+    private Map<Integer, NeuralNetConfiguration.Builder> layerwise;
+    private NeuralNetConfiguration.Builder globalConfig;
+
+    // Constructor
+    public ListBuilder(NeuralNetConfiguration.Builder globalConfig, Map<Integer, NeuralNetConfiguration.Builder> layerMap) {
+        super();
+        this.globalConfig = globalConfig;
+        this.layerwise = layerMap;
+    }
+
+    public ListBuilder(NeuralNetConfiguration.Builder globalConfig) {
+        this(globalConfig, new HashMap<>());
+    }
+
+    public ListBuilder layer(int ind, @NonNull Layer layer) {
+        if (layerwise.containsKey(ind)) {
+            log.info("Layer index {} already exists, layer of type {} will be replace by layer type {}",
+                    ind, layerwise.get(ind).getClass().getSimpleName(), layer.getClass().getSimpleName());
+            layerwise.get(ind).layer(layer);
+        } else {
+            layerwise.put(ind, globalConfig.clone().layer(layer));
+        }
+        if (layerCounter < ind) {
+            //Edge case: user is mixing .layer(Layer) and .layer(int, Layer) calls
+            //This should allow a .layer(A, X) and .layer(Y) to work such that layer Y is index (A+1)
+            layerCounter = ind;
+        }
+        return this;
+    }
+
+    public ListBuilder layer(Layer layer) {
+        return layer(++layerCounter, layer);
+    }
+
+    public Map<Integer, NeuralNetConfiguration.Builder> getLayerwise() {
+        return layerwise;
+    }
+
+    @Override
+    public ListBuilder overrideNinUponBuild(boolean overrideNinUponBuild) {
+        super.overrideNinUponBuild(overrideNinUponBuild);
+        return this;
+    }
+
+    @Override
+    public ListBuilder inputPreProcessor(Integer layer, InputPreProcessor processor) {
+        super.inputPreProcessor(layer, processor);
+        return this;
+    }
+
+
+
+    @Override
+    public ListBuilder cacheMode(@NonNull CacheMode cacheMode) {
+        super.cacheMode(cacheMode);
+        return this;
+    }
+
+
+
+    @Override
+    public ListBuilder tBPTTLength(int bpttLength) {
+        super.tBPTTLength(bpttLength);
+        return this;
+    }
+
+    @Override
+    public ListBuilder tBPTTForwardLength(int forwardLength) {
+        super.tBPTTForwardLength(forwardLength);
+        return this;
+    }
+
+    @Override
+    public ListBuilder tBPTTBackwardLength(int backwardLength) {
+        super.tBPTTBackwardLength(backwardLength);
+        return this;
+    }
+
+
+    @Override
+    public ListBuilder validateOutputLayerConfig(boolean validate) {
+        super.validateOutputLayerConfig(validate);
+        return this;
+    }
+
+    @Override
+    public ListBuilder validateTbpttConfig(boolean validate) {
+        super.validateTbpttConfig(validate);
+        return this;
+    }
+
+    @Override
+    public ListBuilder dataType(@NonNull DataType dataType) {
+        super.dataType(dataType);
+        return this;
+    }
+
+    @Override
+    protected void finalize() throws Throwable {
+        super.finalize();
+    }
+
+    @Override
+    public ListBuilder setInputType(InputType inputType) {
+        return (ListBuilder) super.setInputType(inputType);
+    }
+
+    /**
+     * A convenience method for setting input types: note that for example .inputType().convolutional(h,w,d)
+     * is equivalent to .setInputType(InputType.convolutional(h,w,d))
+     */
+    public InputTypeBuilder inputType() {
+        return new InputTypeBuilder();
+    }
+
+    /**
+     * For the (perhaps partially constructed) network configuration, return a list of activation sizes for each
+     * layer in the network.<br>
+     * Note: To use this method, the network input type must have been set using {@link #setInputType(InputType)} first
+     *
+     * @return A list of activation types for the network, indexed by layer number
+     */
+    public List<InputType> getLayerActivationTypes() {
+        Preconditions.checkState(inputType != null, "Can only calculate activation types if input type has" +
+                "been set. Use setInputType(InputType)");
+
+        MultiLayerConfiguration conf;
+        try {
+            conf = build();
+        } catch (Exception e) {
+            throw new RuntimeException("Error calculating layer activation types: error instantiating MultiLayerConfiguration", e);
+        }
+
+        return conf.getLayerActivationTypes(inputType);
+    }
+
+    /**
+     * Build the multi layer network
+     * based on this neural network and
+     * overr ridden parameters
+     *
+     * @return the configuration to build
+     */
+    public MultiLayerConfiguration build() {
+        List<NeuralNetConfiguration> list = new ArrayList<>();
+        if (layerwise.isEmpty())
+            throw new IllegalStateException("Invalid configuration: no layers defined");
+        for (int i = 0; i < layerwise.size(); i++) {
+            if (layerwise.get(i) == null) {
+                throw new IllegalStateException("Invalid configuration: layer number " + i
+                        + " not specified. Expect layer " + "numbers to be 0 to " + (layerwise.size() - 1)
+                        + " inclusive (number of layers defined: " + layerwise.size() + ")");
+            }
+            if (layerwise.get(i).getLayer() == null)
+                throw new IllegalStateException("Cannot construct network: Layer config for" + "layer with index "
+                        + i + " is not defined)");
+
+            //Layer names: set to default, if not set
+            if (layerwise.get(i).getLayer().getLayerName() == null) {
+                layerwise.get(i).getLayer().setLayerName("layer" + i);
+            }
+
+            list.add(layerwise.get(i).build());
+        }
+
+        WorkspaceMode wsmTrain = (globalConfig.setTWM ? globalConfig.trainingWorkspaceMode : trainingWorkspaceMode);
+        WorkspaceMode wsmTest = (globalConfig.setIWM ? globalConfig.inferenceWorkspaceMode : inferenceWorkspaceMode);
+
+
+        MultiLayerConfiguration.Builder builder = new MultiLayerConfiguration.Builder().inputPreProcessors(inputPreProcessors)
+                .backpropType(backpropType).tBPTTForwardLength(tbpttFwdLength)
+                .tBPTTBackwardLength(tbpttBackLength).setInputType(this.inputType)
+                .trainingWorkspaceMode(wsmTrain).cacheMode(globalConfig.cacheMode)
+                .inferenceWorkspaceMode(wsmTest).confs(list).validateOutputLayerConfig(validateOutputConfig)
+                .overrideNinUponBuild(overrideNinUponBuild)
+                .dataType(globalConfig.dataType);
+        return builder.build();
+    }
+
+    /**
+     * Helper class for setting input types
+     */
+    public class InputTypeBuilder {
+        /**
+         * See {@link InputType#convolutional(long, long, long)}
+         */
+        public ListBuilder convolutional(int height, int width, int depth) {
+            return ListBuilder.this.setInputType(InputType.convolutional(height, width, depth));
+        }
+
+        /**
+         * * See {@link InputType#convolutionalFlat(long, long, long)}
+         */
+        public ListBuilder convolutionalFlat(int height, int width, int depth) {
+            return ListBuilder.this.setInputType(InputType.convolutionalFlat(height, width, depth));
+        }
+
+        /**
+         * See {@link InputType#feedForward(long)}
+         */
+        public ListBuilder feedForward(int size) {
+            return ListBuilder.this.setInputType(InputType.feedForward(size));
+        }
+
+        /**
+         * See {@link InputType#recurrent(long)}}
+         */
+        public ListBuilder recurrent(int size) {
+            return ListBuilder.this.setInputType(InputType.recurrent(size));
+        }
+    }
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/conf/MultiLayerConfiguration.java

@@ -0,0 +1,647 @@
+
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.nn.conf;
+
+import lombok.AccessLevel;
+import lombok.AllArgsConstructor;
+import lombok.Data;
+import lombok.Getter;
+import lombok.NoArgsConstructor;
+import lombok.NonNull;
+import lombok.Setter;
+import lombok.extern.slf4j.Slf4j;
+import lombok.val;
+
+import org.deeplearning4j.nn.conf.distribution.Distribution;
+import org.deeplearning4j.nn.conf.inputs.InputType;
+import org.deeplearning4j.nn.conf.layers.*;
+import org.deeplearning4j.nn.conf.layers.recurrent.LastTimeStep;
+import org.deeplearning4j.nn.conf.memory.LayerMemoryReport;
+import org.deeplearning4j.nn.conf.memory.MemoryReport;
+import org.deeplearning4j.nn.conf.memory.NetworkMemoryReport;
+import org.deeplearning4j.nn.conf.serde.ComputationGraphConfigurationDeserializer;
+import org.deeplearning4j.nn.conf.serde.JsonMappers;
+import org.deeplearning4j.nn.conf.serde.MultiLayerConfigurationDeserializer;
+import org.deeplearning4j.nn.weights.IWeightInit;
+import org.deeplearning4j.nn.weights.WeightInit;
+import org.deeplearning4j.util.OutputLayerUtil;
+import org.nd4j.linalg.activations.Activation;
+import org.nd4j.linalg.activations.IActivation;
+import org.nd4j.linalg.api.buffer.DataType;
+import org.nd4j.linalg.factory.Nd4j;
+import org.nd4j.linalg.lossfunctions.LossFunctions;
+import org.nd4j.linalg.lossfunctions.impl.LossBinaryXENT;
+import org.nd4j.linalg.lossfunctions.impl.LossMCXENT;
+import org.nd4j.linalg.lossfunctions.impl.LossMSE;
+import org.nd4j.linalg.lossfunctions.impl.LossNegativeLogLikelihood;
+import org.nd4j.shade.jackson.databind.*;
+import org.nd4j.shade.jackson.databind.deser.BeanDeserializerModifier;
+import org.nd4j.shade.jackson.databind.exc.InvalidTypeIdException;
+import org.nd4j.shade.jackson.databind.module.SimpleModule;
+import org.nd4j.shade.jackson.databind.node.ArrayNode;
+import org.nd4j.shade.jackson.dataformat.yaml.YAMLFactory;
+
+import java.io.IOException;
+import java.io.Serializable;
+import java.util.*;
+
+@Data
+@AllArgsConstructor(access = AccessLevel.PRIVATE)
+@NoArgsConstructor
+@Slf4j
+public class MultiLayerConfiguration implements Serializable, Cloneable {
+
+    protected List<NeuralNetConfiguration> confs;
+    protected Map<Integer, InputPreProcessor> inputPreProcessors = new HashMap<>();
+    protected BackpropType backpropType = BackpropType.Standard;
+    protected int tbpttFwdLength = 20;
+    protected int tbpttBackLength = 20;
+    protected boolean validateOutputLayerConfig = true; //Default to legacy for pre 1.0.0-beta3 networks on deserialization
+
+    @Getter
+    @Setter
+    protected WorkspaceMode trainingWorkspaceMode = WorkspaceMode.ENABLED;
+
+    @Getter
+    @Setter
+    protected WorkspaceMode inferenceWorkspaceMode = WorkspaceMode.ENABLED;
+
+    @Getter
+    @Setter
+    protected CacheMode cacheMode;
+
+    @Getter
+    @Setter
+    protected DataType dataType = DataType.FLOAT;   //Default to float for deserialization of beta3 and earlier nets
+
+    //Counter for the number of parameter updates so far
+    // This is important for learning rate schedules, for example, and is stored here to ensure it is persisted
+    // for Spark and model serialization
+    protected int iterationCount = 0;
+
+    //Counter for the number of epochs completed so far. Used for per-epoch schedules
+    protected int epochCount = 0;
+    private static ObjectMapper mapper = mapper();
+    private static ObjectMapper mapperYaml = mapperYaml();
+
+
+
+    public static ObjectMapper mapperYaml() {
+        ObjectMapper ret = new ObjectMapper(new YAMLFactory());
+        ret.configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false);
+        ret.configure(SerializationFeature.FAIL_ON_EMPTY_BEANS, false);
+        ret.configure(MapperFeature.SORT_PROPERTIES_ALPHABETICALLY, true);
+        ret.enable(SerializationFeature.INDENT_OUTPUT);
+
+        SimpleModule customDeserializerModule = new SimpleModule();
+        customDeserializerModule.setDeserializerModifier(new BeanDeserializerModifier() {
+            @Override
+            public JsonDeserializer<?> modifyDeserializer(DeserializationConfig config, BeanDescription beanDesc,
+                                                          JsonDeserializer<?> deserializer) {
+                //Use our custom deserializers to handle backward compatibility for updaters -> IUpdater
+                if (beanDesc.getBeanClass().equals(MultiLayerConfiguration.class)) {
+                    return new MultiLayerConfigurationDeserializer(deserializer);
+                }
+                return deserializer;
+            }
+        });
+
+        ret.registerModule(customDeserializerModule);
+        return ret;
+    }
+
+
+    public static ObjectMapper mapper() {
+        ObjectMapper ret = new ObjectMapper();
+        ret.configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false);
+        ret.configure(SerializationFeature.FAIL_ON_EMPTY_BEANS, false);
+        ret.configure(MapperFeature.SORT_PROPERTIES_ALPHABETICALLY, true);
+        ret.enable(SerializationFeature.INDENT_OUTPUT);
+        SimpleModule customDeserializerModule = new SimpleModule();
+        customDeserializerModule.setDeserializerModifier(new BeanDeserializerModifier() {
+            @Override
+            public JsonDeserializer<?> modifyDeserializer(DeserializationConfig config, BeanDescription beanDesc,
+                                                          JsonDeserializer<?> deserializer) {
+                //Use our custom deserializers to handle backward compatibility for updaters -> IUpdater
+                if (beanDesc.getBeanClass().equals(MultiLayerConfiguration.class)) {
+                    return new MultiLayerConfigurationDeserializer(deserializer);
+                }
+                return deserializer;
+            }
+        });
+
+        ret.registerModule(customDeserializerModule);
+        return ret;
+    }
+
+    public int getEpochCount() {
+        return epochCount;
+    }
+
+    public void setEpochCount(int epochCount) {
+        this.epochCount = epochCount;
+        for (int i = 0; i < confs.size(); i++) {
+            getConf(i).setEpochCount(epochCount);
+        }
+    }
+
+    /**
+     * @return JSON representation of NN configuration
+     */
+    public String toYaml() {
+        try {
+            return mapperYaml.writeValueAsString(this);
+        } catch (org.nd4j.shade.jackson.core.JsonProcessingException e) {
+            throw new RuntimeException(e);
+        }
+
+    }
+
+    /**
+     * Create a neural net configuration from json
+     *
+     * @param json the neural net configuration from json
+     * @return {@link MultiLayerConfiguration}
+     */
+    public static MultiLayerConfiguration fromYaml(String json) {
+        try {
+            return mapperYaml.readValue(json, MultiLayerConfiguration.class);
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+
+    /**
+     * @return JSON representation of NN configuration
+     */
+    public String toJson() {
+        //JSON mappers are supposed to be thread safe: however, in practice they seem to miss fields occasionally
+        //when writeValueAsString is used by multiple threads. This results in invalid JSON. See issue #3243
+        try {
+            return mapper.writeValueAsString(this);
+        } catch (org.nd4j.shade.jackson.core.JsonProcessingException e) {
+            throw new RuntimeException(e);
+        }
+
+    }
+
+    /**
+     * Create a neural net configuration from json
+     *
+     * @param json the neural net configuration from json
+     * @return {@link MultiLayerConfiguration}
+     */
+    public static  MultiLayerConfiguration fromJson(String json) {
+          ObjectMapper mapper1 = mapper();
+        MultiLayerConfiguration conf;
+        try {
+            conf = mapper1.readValue(json, MultiLayerConfiguration.class);
+        } catch (InvalidTypeIdException e){
+            if(e.getMessage().contains("@class")) {
+                try {
+                    //JSON may be legacy (1.0.0-alpha or earlier), attempt to load it using old format
+                    return JsonMappers.getLegacyMapper().readValue(json, MultiLayerConfiguration.class);
+                } catch (InvalidTypeIdException e2) {
+                    //Check for legacy custom layers: "Could not resolve type id 'CustomLayer' as a subtype of [simple type, class org.deeplearning4j.nn.conf.layers.Layer]: known type ids = [Bidirectional, CenterLossOutputLayer, CnnLossLayer, ..."
+                    //1.0.0-beta5: dropping support for custom layers defined in pre-1.0.0-beta format. Built-in layers from these formats still work
+                    String msg = e2.getMessage();
+                    if(msg != null && msg.contains("Could not resolve type id")){
+                        throw new RuntimeException("Error deserializing MultiLayerConfiguration - configuration may have a custom " +
+                                "layer, vertex or preprocessor, in pre version 1.0.0-beta JSON format.\nModels in legacy format with custom" +
+                                " layers should be loaded in 1.0.0-beta to 1.0.0-beta4 and saved again, before loading in the current version of DL4J", e);
+                    }
+                    throw new RuntimeException(e2);
+                } catch (IOException e2) {
+                    throw new RuntimeException(e2);
+                }
+            }
+            throw new RuntimeException(e);
+        } catch (IOException e) {
+            //Check if this exception came from legacy deserializer...
+            String msg = e.getMessage();
+            if (msg != null && msg.contains("legacy")) {
+                throw new RuntimeException("Error deserializing MultiLayerConfiguration - configuration may have a custom " +
+                        "layer, vertex or preprocessor, in pre version 1.0.0-alpha JSON format. These layers can be " +
+                        "deserialized by first registering them with NeuralNetConfiguration.registerLegacyCustomClassesForJSON(Class...)", e);
+            }
+            throw new RuntimeException(e);
+        }
+
+
+        //To maintain backward compatibility after loss function refactoring (configs generated with v0.5.0 or earlier)
+        // Previously: enumeration used for loss functions. Now: use classes
+        // IN the past, could have only been an OutputLayer or RnnOutputLayer using these enums
+        int layerCount = 0;
+        JsonNode confs = null;
+        for (NeuralNetConfiguration nnc : conf.getConfs()) {
+            Layer l = nnc.getLayer();
+            if (l instanceof BaseOutputLayer && ((BaseOutputLayer) l).getLossFn() == null) {
+                //lossFn field null -> may be an old config format, with lossFunction field being for the enum
+                //if so, try walking the JSON graph to extract out the appropriate enum value
+
+                BaseOutputLayer ol = (BaseOutputLayer) l;
+                try {
+                    JsonNode jsonNode = mapper.readTree(json);
+                    if (confs == null) {
+                        confs = jsonNode.get("confs");
+                    }
+                    if (confs instanceof ArrayNode) {
+                        ArrayNode layerConfs = (ArrayNode) confs;
+                        JsonNode outputLayerNNCNode = layerConfs.get(layerCount);
+                        if (outputLayerNNCNode == null)
+                            return conf; //Should never happen...
+                        JsonNode outputLayerNode = outputLayerNNCNode.get("layer");
+
+                        JsonNode lossFunctionNode = null;
+                        if (outputLayerNode.has("output")) {
+                            lossFunctionNode = outputLayerNode.get("output").get("lossFunction");
+                        } else if (outputLayerNode.has("rnnoutput")) {
+                            lossFunctionNode = outputLayerNode.get("rnnoutput").get("lossFunction");
+                        }
+
+                        if (lossFunctionNode != null) {
+                            String lossFunctionEnumStr = lossFunctionNode.asText();
+                            LossFunctions.LossFunction lossFunction = null;
+                            try {
+                                lossFunction = LossFunctions.LossFunction.valueOf(lossFunctionEnumStr);
+                            } catch (Exception e) {
+                                log.warn("OutputLayer with null LossFunction or pre-0.6.0 loss function configuration detected: could not parse JSON",
+                                        e);
+                            }
+
+                            if (lossFunction != null) {
+                                switch (lossFunction) {
+                                    case MSE:
+                                        ol.setLossFn(new LossMSE());
+                                        break;
+                                    case XENT:
+                                        ol.setLossFn(new LossBinaryXENT());
+                                        break;
+                                    case NEGATIVELOGLIKELIHOOD:
+                                        ol.setLossFn(new LossNegativeLogLikelihood());
+                                        break;
+                                    case MCXENT:
+                                        ol.setLossFn(new LossMCXENT());
+                                        break;
+
+                                    //Remaining: TODO
+                                    case SQUARED_LOSS:
+                                    case RECONSTRUCTION_CROSSENTROPY:
+                                    default:
+                                        log.warn("OutputLayer with null LossFunction or pre-0.6.0 loss function configuration detected: could not set loss function for {}",
+                                                lossFunction);
+                                        break;
+                                }
+                            }
+                        }
+
+                    } else {
+                        log.warn("OutputLayer with null LossFunction or pre-0.6.0 loss function configuration detected: could not parse JSON: layer 'confs' field is not an ArrayNode (is: {})",
+                                (confs != null ? confs.getClass() : null));
+                    }
+                } catch (IOException e) {
+                    log.warn("OutputLayer with null LossFunction or pre-0.6.0 loss function configuration detected: could not parse JSON",
+                            e);
+                    break;
+                }
+            }
+
+            //Also, pre 0.7.2: activation functions were Strings ("activationFunction" field), not classes ("activationFn")
+            //Try to load the old format if necessary, and create the appropriate IActivation instance
+            if ((l instanceof BaseLayer) && ((BaseLayer) l).getActivationFn() == null) {
+                try {
+                    JsonNode jsonNode = mapper.readTree(json);
+                    if (confs == null) {
+                        confs = jsonNode.get("confs");
+                    }
+                    if (confs instanceof ArrayNode) {
+                        ArrayNode layerConfs = (ArrayNode) confs;
+                        JsonNode outputLayerNNCNode = layerConfs.get(layerCount);
+                        if (outputLayerNNCNode == null)
+                            return conf; //Should never happen...
+                        JsonNode layerWrapperNode = outputLayerNNCNode.get("layer");
+
+                        if (layerWrapperNode == null || layerWrapperNode.size() != 1) {
+                            continue;
+                        }
+
+                        JsonNode layerNode = layerWrapperNode.elements().next();
+                        JsonNode activationFunction = layerNode.get("activationFunction"); //Should only have 1 element: "dense", "output", etc
+
+                        if (activationFunction != null) {
+                            IActivation ia = Activation.fromString(activationFunction.asText()).getActivationFunction();
+                            ((BaseLayer) l).setActivationFn(ia);
+                        }
+                    }
+
+                } catch (IOException e) {
+                    log.warn("Layer with null ActivationFn field or pre-0.7.2 activation function detected: could not parse JSON",
+                            e);
+                }
+            }
+
+            if(!handleLegacyWeightInitFromJson(json, l, mapper, confs, layerCount)) {
+                return conf;
+            }
+
+            layerCount++;
+        }
+        return conf;
+    }
+
+    /**
+     * Handle {@link WeightInit} and {@link Distribution} from legacy configs in Json format. Copied from handling of {@link Activation}
+     * above.
+     * @return True if all is well and layer iteration shall continue. False else-wise.
+     */
+    private static boolean handleLegacyWeightInitFromJson(String json, Layer l, ObjectMapper mapper, JsonNode confs, int layerCount) {
+        if ((l instanceof BaseLayer) && ((BaseLayer) l).getWeightInitFn() == null) {
+            try {
+                JsonNode jsonNode = mapper.readTree(json);
+                if (confs == null) {
+                    confs = jsonNode.get("confs");
+                }
+                if (confs instanceof ArrayNode) {
+                    ArrayNode layerConfs = (ArrayNode) confs;
+                    JsonNode outputLayerNNCNode = layerConfs.get(layerCount);
+                    if (outputLayerNNCNode == null)
+                        return false; //Should never happen...
+                    JsonNode layerWrapperNode = outputLayerNNCNode.get("layer");
+
+                    if (layerWrapperNode == null || layerWrapperNode.size() != 1) {
+                        return true;
+                    }
+
+                    JsonNode layerNode = layerWrapperNode.elements().next();
+                    JsonNode weightInit = layerNode.get("weightInit"); //Should only have 1 element: "dense", "output", etc
+                    JsonNode distribution = layerNode.get("dist");
+
+                    Distribution dist = null;
+                    if(distribution != null) {
+                        dist = mapper.treeToValue(distribution, Distribution.class);
+                    }
+
+                    if (weightInit != null) {
+                        IWeightInit wi = WeightInit.valueOf(weightInit.asText()).getWeightInitFunction(dist);
+                        ((BaseLayer) l).setWeightInitFn(wi);
+                    }
+                }
+
+            } catch (IOException e) {
+                log.warn("Layer with null WeightInit detected: " + l.getLayerName() + ", could not parse JSON",
+                        e);
+            }
+        }
+        return true;
+
+    }
+
+    @Override
+    public String toString() {
+        return toJson();
+    }
+
+    public NeuralNetConfiguration getConf(int i) {
+        return confs.get(i);
+    }
+
+    @Override
+    public MultiLayerConfiguration clone() {
+        try {
+            MultiLayerConfiguration clone = (MultiLayerConfiguration) super.clone();
+
+            if (clone.confs != null) {
+                List<NeuralNetConfiguration> list = new ArrayList<>();
+                for (NeuralNetConfiguration conf : clone.confs) {
+                    list.add(conf.clone());
+                }
+                clone.confs = list;
+            }
+
+            if (clone.inputPreProcessors != null) {
+                Map<Integer, InputPreProcessor> map = new HashMap<>();
+                for (Map.Entry<Integer, InputPreProcessor> entry : clone.inputPreProcessors.entrySet()) {
+                    map.put(entry.getKey(), entry.getValue().clone());
+                }
+                clone.inputPreProcessors = map;
+            }
+
+            clone.inferenceWorkspaceMode = this.inferenceWorkspaceMode;
+            clone.trainingWorkspaceMode = this.trainingWorkspaceMode;
+            clone.cacheMode = this.cacheMode;
+            clone.validateOutputLayerConfig = this.validateOutputLayerConfig;
+            clone.dataType = this.dataType;
+
+            return clone;
+
+        } catch (CloneNotSupportedException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    public InputPreProcessor getInputPreProcess(int curr) {
+        return inputPreProcessors.get(curr);
+    }
+
+    /**
+     * Get a {@link MemoryReport} for the given MultiLayerConfiguration. This is used to estimate the
+     * memory requirements for the given network configuration and input
+     *
+     * @param inputType Input types for the network
+     * @return Memory report for the network
+     */
+    public NetworkMemoryReport getMemoryReport(InputType inputType) {
+
+        Map<String, MemoryReport> memoryReportMap = new LinkedHashMap<>();
+        int nLayers = confs.size();
+        for (int i = 0; i < nLayers; i++) {
+            String layerName = confs.get(i).getLayer().getLayerName();
+            if (layerName == null) {
+                layerName = String.valueOf(i);
+            }
+
+            //Pass input type through preprocessor, if necessary
+            InputPreProcessor preproc = getInputPreProcess(i);
+            //TODO memory requirements for preprocessor
+            if (preproc != null) {
+                inputType = preproc.getOutputType(inputType);
+            }
+
+            LayerMemoryReport report = confs.get(i).getLayer().getMemoryReport(inputType);
+            memoryReportMap.put(layerName, report);
+
+            inputType = confs.get(i).getLayer().getOutputType(i, inputType);
+        }
+
+        return new NetworkMemoryReport(memoryReportMap, MultiLayerConfiguration.class, "MultiLayerNetwork", inputType);
+    }
+
+    /**
+     * For the given input shape/type for the network, return a list of activation sizes for each layer in the network.<br>
+     * i.e., list.get(i) is the output activation sizes for layer i
+     *
+     * @param inputType Input type for the network
+     * @return A lits of activation types for the network, indexed by layer number
+     */
+    public List<InputType> getLayerActivationTypes(@NonNull InputType inputType) {
+        List<InputType> out = new ArrayList<>();
+        int nLayers = confs.size();
+        for (int i = 0; i < nLayers; i++) {
+            InputPreProcessor preproc = getInputPreProcess(i);
+            if (preproc != null) {
+                inputType = preproc.getOutputType(inputType);
+            }
+
+            inputType = confs.get(i).getLayer().getOutputType(i, inputType);
+            out.add(inputType);
+        }
+        return out;
+    }
+
+    @Data
+    public static class Builder extends BaseBuilder {
+
+        public MultiLayerConfiguration build() {
+            //Validate BackpropType setting
+            if ((tbpttBackLength != DEFAULT_TBPTT_LENGTH || tbpttFwdLength != DEFAULT_TBPTT_LENGTH) && backpropType != BackpropType.TruncatedBPTT) {
+                log.warn("Truncated backpropagation through time lengths have been configured with values " + tbpttFwdLength
+                        + " and " + tbpttBackLength + " but backprop type is set to " + backpropType + ". TBPTT configuration" +
+                        " settings will only take effect if backprop type is set to BackpropType.TruncatedBPTT");
+            }
+
+            if(backpropType == BackpropType.TruncatedBPTT && validateTbpttConfig) {
+                //Check for invalid combination - tbptt plus LastTimeStepLayer or
+                for( int i = 0; i < confs.size(); i++) {
+                    Layer l = confs.get(i).getLayer();
+                    if(l instanceof LastTimeStep || l instanceof GlobalPoolingLayer) {
+                        throw new IllegalStateException("Invalid network configuration detected: Truncated backpropagation through time (TBPTT)" +
+                                " cannot be used with layer " + i + " of type " + l.getClass().getName() + ": TBPTT is incompatible with this layer type (which is designed " +
+                                "to process entire sequences at once, and does support the type of sequence segments that TPBTT uses).\n" +
+                                "This check can be disabled using validateTbpttConfig(false) but this is not recommended.");
+                    }
+                }
+            }
+
+
+            if (inputType == null && inputPreProcessors.get(0) == null) {
+                //User hasn't set the InputType. Sometimes we can infer it...
+                // For example, Dense/RNN layers, where preprocessor isn't set -> user is *probably* going to feed in
+                // standard feedforward or RNN data
+                //This isn't the most elegant implementation, but should avoid breaking backward compatibility here
+                //Can't infer InputType for CNN layers, however (don't know image dimensions/depth)
+                Layer firstLayer = confs.get(0).getLayer();
+                if (firstLayer instanceof BaseRecurrentLayer) {
+                    BaseRecurrentLayer brl = (BaseRecurrentLayer) firstLayer;
+                    val nIn = brl.getNIn();
+                    if (nIn > 0) {
+                        inputType = InputType.recurrent(nIn, brl.getRnnDataFormat());
+                    }
+                } else if (firstLayer instanceof DenseLayer || firstLayer instanceof EmbeddingLayer
+                        || firstLayer instanceof OutputLayer) {
+                    //Can't just use "instanceof FeedForwardLayer" here. ConvolutionLayer is also a FeedForwardLayer
+                    FeedForwardLayer ffl = (FeedForwardLayer) firstLayer;
+                    val nIn = ffl.getNIn();
+                    if (nIn > 0) {
+                        inputType = InputType.feedForward(nIn);
+                    }
+                }
+            }
+
+
+            //Add preprocessors and set nIns, if InputType has been set
+            // Builder.inputType field can be set in 1 of 4 ways:
+            // 1. User calls setInputType directly
+            // 2. Via ConvolutionLayerSetup -> internally calls setInputType(InputType.convolutional(...))
+            // 3. Via the above code: i.e., assume input is as expected  by the RNN or dense layer -> sets the inputType field
+            if (inputType != null) {
+                InputType currentInputType = inputType;
+                for (int i = 0; i < confs.size(); i++) {
+                    Layer l = confs.get(i).getLayer();
+                    if (inputPreProcessors.get(i) == null) {
+                        //Don't override preprocessor setting, but set preprocessor if required...
+                        InputPreProcessor inputPreProcessor = l.getPreProcessorForInputType(currentInputType);
+                        if (inputPreProcessor != null) {
+                            inputPreProcessors.put(i, inputPreProcessor);
+                        }
+                    }
+
+                    InputPreProcessor inputPreProcessor = inputPreProcessors.get(i);
+                    if (inputPreProcessor != null) {
+                        currentInputType = inputPreProcessor.getOutputType(currentInputType);
+                    }
+                    if(i > 0) {
+                        Layer layer = confs.get(i - 1).getLayer();
+                        //convolution 1d is an edge case where it has rnn input type but the filters
+                        //should be the output
+                        if(layer instanceof Convolution1DLayer) {
+                            if(l instanceof DenseLayer && inputType instanceof InputType.InputTypeRecurrent) {
+                                FeedForwardLayer feedForwardLayer = (FeedForwardLayer) l;
+                                if(inputType instanceof InputType.InputTypeRecurrent) {
+                                    InputType.InputTypeRecurrent recurrent = (InputType.InputTypeRecurrent) inputType;
+                                    feedForwardLayer.setNIn(recurrent.getTimeSeriesLength());
+                                }
+                            }
+                            else
+                                l.setNIn(currentInputType, overrideNinUponBuild); //Don't override the nIn setting, if it's manually set by the user
+                        }
+                        else
+                            l.setNIn(currentInputType, overrideNinUponBuild); //Don't override the nIn setting, if it's manually set by the user
+
+                    }
+                    else
+                        l.setNIn(currentInputType, overrideNinUponBuild); //Don't override the nIn setting, if it's manually set by the user
+
+
+                    currentInputType = l.getOutputType(i, currentInputType);
+                }
+
+            }
+
+            MultiLayerConfiguration conf = new MultiLayerConfiguration();
+            conf.confs = this.confs;
+            conf.inputPreProcessors = inputPreProcessors;
+            conf.backpropType = backpropType;
+            conf.tbpttFwdLength = tbpttFwdLength;
+            conf.tbpttBackLength = tbpttBackLength;
+            conf.trainingWorkspaceMode = trainingWorkspaceMode;
+            conf.inferenceWorkspaceMode = inferenceWorkspaceMode;
+            conf.cacheMode = cacheMode;
+            conf.dataType = dataType;
+
+            Nd4j.getRandom().setSeed(conf.getConf(0).getSeed());
+
+            //Validate output layer configuration
+            if (validateOutputConfig) {
+                //Validate output layer configurations...
+                for (NeuralNetConfiguration n : conf.getConfs()) {
+                    Layer l = n.getLayer();
+                    OutputLayerUtil.validateOutputLayer(l.getLayerName(), l); //No-op for non output/loss layers
+                }
+            }
+
+            return conf;
+
+        }
+    }
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/conf/RNNFormat.java

@@ -0,0 +1,27 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.nn.conf;
+
+public enum RNNFormat implements DataFormat {
+    NCW,
+    NWC
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/conf/Updater.java

@@ -0,0 +1,60 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.nn.conf;
+
+import org.nd4j.linalg.learning.config.*;
+
+/**
+ *
+ * All the possible different updaters
+ *
+ * @author Adam Gibson
+ */
+public enum Updater {
+    SGD, ADAM, ADAMAX, ADADELTA, NESTEROVS, NADAM, ADAGRAD, RMSPROP, NONE, @Deprecated CUSTOM;
+
+
+    public IUpdater getIUpdaterWithDefaultConfig() {
+        switch (this) {
+            case SGD:
+                return new Sgd();
+            case ADAM:
+                return new Adam();
+            case ADAMAX:
+                return new AdaMax();
+            case ADADELTA:
+                return new AdaDelta();
+            case NESTEROVS:
+                return new Nesterovs();
+            case NADAM:
+                return new Nadam();
+            case ADAGRAD:
+                return new AdaGrad();
+            case RMSPROP:
+                return new RmsProp();
+            case NONE:
+                return new NoOp();
+            case CUSTOM:
+            default:
+                throw new UnsupportedOperationException("Unknown or not supported updater: " + this);
+        }
+    }
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/conf/WorkspaceMode.java

@@ -0,0 +1,36 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.nn.conf;
+
+public enum WorkspaceMode {
+    NONE, // workspace won't be used
+    ENABLED,
+    /**
+     * @deprecated Use {@link #ENABLED} instead
+     */
+    @Deprecated
+    SINGLE, // one external workspace
+    /**
+     * @deprecated Use {@link #ENABLED} instead
+     */
+    @Deprecated
+    SEPARATE, // one external workspace, one FF workspace, one BP workspace <-- default one
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/conf/constraint/BaseConstraint.java

@@ -0,0 +1,94 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.nn.conf.constraint;
+
+import lombok.AllArgsConstructor;
+import lombok.Data;
+import lombok.EqualsAndHashCode;
+import lombok.Getter;
+import lombok.Setter;
+
+import org.apache.commons.lang3.ArrayUtils;
+import org.deeplearning4j.nn.api.Layer;
+import org.deeplearning4j.nn.api.ParamInitializer;
+import org.deeplearning4j.nn.api.layers.LayerConstraint;
+import org.nd4j.linalg.api.ndarray.INDArray;
+
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Set;
+
+
+@AllArgsConstructor
+@EqualsAndHashCode
+@Data
+public abstract class BaseConstraint implements LayerConstraint {
+    public static final double DEFAULT_EPSILON = 1e-6;
+    @Setter
+    @Getter
+    protected Set<String> params = new HashSet<>();
+    protected double epsilon = 1e-6;
+    protected long[] dimensions;
+
+    protected BaseConstraint(){
+        //No arg for json ser/de
+    }
+
+    protected BaseConstraint(Set<String> paramNames, long... dimensions){
+        this(paramNames, DEFAULT_EPSILON, dimensions);
+    }
+
+    @Override
+    public void applyConstraint(Layer layer, int iteration, int epoch) {
+        Map<String,INDArray> paramTable = layer.paramTable();
+        if(paramTable == null || paramTable.isEmpty() ){
+            return;
+        }
+
+        ParamInitializer i = layer.conf().getLayer().initializer();
+        for(Map.Entry<String,INDArray> e : paramTable.entrySet()){
+            if(params.contains(e.getKey())){
+                apply(e.getValue());
+            }
+            if (params != null && params.contains(e.getKey())) {
+                apply(e.getValue());
+            }
+        }
+    }
+
+    public abstract void apply(INDArray param);
+
+    public abstract BaseConstraint clone();
+
+    public static long[] getBroadcastDims(long[] reduceDimensions, int rank) {
+        long[] out = new long[rank - reduceDimensions.length];
+        if(rank < 1 || reduceDimensions.length < 1 || out.length < 1) {
+            return new long[]{0};
+        }
+        int outPos = 0;
+        for( int i = 0; i < rank; i++) {
+            if(!ArrayUtils.contains(reduceDimensions, i)) {
+                out[outPos++] = i;
+            }
+        }
+        return out;
+    }
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/conf/constraint/MaxNormConstraint.java

@@ -0,0 +1,83 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.nn.conf.constraint;
+
+import lombok.Data;
+import lombok.EqualsAndHashCode;
+import org.nd4j.linalg.api.ndarray.INDArray;
+import org.nd4j.linalg.factory.Broadcast;
+import org.nd4j.linalg.indexing.BooleanIndexing;
+import org.nd4j.linalg.indexing.conditions.Conditions;
+
+import java.util.Collections;
+import java.util.Set;
+
+@Data
+@EqualsAndHashCode(callSuper = true)
+public class MaxNormConstraint extends BaseConstraint {
+
+    private double maxNorm;
+
+    private MaxNormConstraint(){
+        //No arg for json ser/de
+    }
+
+    /**
+     * @param maxNorm        Maximum L2 value
+     * @param paramNames     Which parameter names to apply constraint to
+     * @param dimensions     Dimensions to apply to. For DenseLayer, OutputLayer, RnnOutputLayer, LSTM, etc: this should
+     *                       be dimension 1. For CNNs, this should be dimensions [1,2,3] corresponding to last 3 of
+     *                       parameters which have order [depthOut, depthIn, kH, kW]
+     */
+    public MaxNormConstraint(double maxNorm, Set<String> paramNames, long... dimensions){
+        super(paramNames, DEFAULT_EPSILON, dimensions);
+        this.maxNorm = maxNorm;
+    }
+
+    /**
+     * Apply to weights but not biases by default
+     *
+     * @param maxNorm        Maximum L2 value
+     * @param dimensions     Dimensions to apply to. For DenseLayer, OutputLayer, RnnOutputLayer, LSTM, etc: this should
+     *                       be dimension 1. For CNNs, this should be dimensions [1,2,3] corresponding to last 3 of
+     *                       parameters which have order [depthOut, depthIn, kH, kW]
+     */
+    public MaxNormConstraint(double maxNorm, long... dimensions) {
+
+        this(maxNorm, Collections.<String>emptySet(), dimensions);
+    }
+
+    @Override
+    public void apply(INDArray param){
+        INDArray norm = param.norm2(dimensions);
+        INDArray clipped = norm.unsafeDuplication();
+        BooleanIndexing.replaceWhere(clipped, maxNorm, Conditions.greaterThan(maxNorm));
+        norm.addi(epsilon);
+        clipped.divi(norm);
+
+        Broadcast.mul(param, clipped, param, getBroadcastDims(dimensions, param.rank()));
+    }
+
+    @Override
+    public MaxNormConstraint clone() {
+        return new MaxNormConstraint(maxNorm,  params, dimensions);
+    }
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/conf/constraint/MinMaxNormConstraint.java

@@ -0,0 +1,119 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.nn.conf.constraint;
+
+import lombok.Data;
+import lombok.EqualsAndHashCode;
+import org.nd4j.linalg.api.ndarray.INDArray;
+import org.nd4j.linalg.api.ops.CustomOp;
+import org.nd4j.linalg.api.ops.DynamicCustomOp;
+import org.nd4j.linalg.factory.Broadcast;
+import org.nd4j.linalg.factory.Nd4j;
+
+import java.util.Collections;
+import java.util.Set;
+
+@Data
+@EqualsAndHashCode(callSuper = true)
+public class MinMaxNormConstraint extends BaseConstraint {
+    public static final double DEFAULT_RATE = 1.0;
+
+    private double min;
+    private double max;
+    private double rate;
+
+    private MinMaxNormConstraint(){
+        //No arg for json ser/de
+    }
+
+    /**
+     * Apply to weights but not biases by default
+     *
+     * @param max            Maximum L2 value
+     * @param min            Minimum L2 value
+     * @param dimensions     Dimensions to apply to. For DenseLayer, OutputLayer, RnnOutputLayer, LSTM, etc: this should
+     *                       be dimension 1. For CNNs, this should be dimensions [1,2,3] corresponding to last 3 of
+     *                       parameters which have order [depthOut, depthIn, kH, kW]
+     */
+    public MinMaxNormConstraint(double min, double max, long... dimensions){
+        this(min, max, DEFAULT_RATE, null, dimensions);
+    }
+
+    /**
+     * Apply to weights but not biases by default
+     *
+     * @param max            Maximum L2 value
+     * @param min            Minimum L2 value
+     * @param rate           Constraint rate
+     * @param dimensions     Dimensions to apply to. For DenseLayer, OutputLayer, RnnOutputLayer, LSTM, etc: this should
+     *                       be dimension 1. For CNNs, this should be dimensions [1,2,3] corresponding to last 3 of
+     *                       parameters which have order [depthOut, depthIn, kH, kW]
+     */
+    public MinMaxNormConstraint(double min, double max, double rate, long... dimensions){
+        this(min, max, rate, Collections.<String>emptySet(), dimensions);
+    }
+
+    /**
+     *
+     * @param max            Maximum L2 value
+     * @param min            Minimum L2 value
+     * @param rate           Constraint rate
+     * @param paramNames     Which parameter names to apply constraint to
+     * @param dimensions     Dimensions to apply to. For DenseLayer, OutputLayer, RnnOutputLayer, LSTM, etc: this should
+     *                       be dimension 1. For CNNs, this should be dimensions [1,2,3] corresponding to last 3 of
+     *                       parameters which have order [depthOut, depthIn, kH, kW]
+     */
+    public MinMaxNormConstraint(double min, double max, double rate, Set<String> paramNames, long... dimensions){
+        super(paramNames, dimensions);
+        if(rate <= 0 || rate > 1.0){
+            throw new IllegalStateException("Invalid rate: must be in interval (0,1]: got " + rate);
+        }
+        this.min = min;
+        this.max = max;
+        this.rate = rate;
+    }
+
+    @Override
+    public void apply(INDArray param) {
+        INDArray norm = param.norm2(dimensions);
+        INDArray clipped = norm.unsafeDuplication();
+        CustomOp op = DynamicCustomOp.builder("clipbyvalue")
+                .addInputs(clipped)
+                .callInplace(true)
+                .addFloatingPointArguments(min, max)
+                .build();
+        Nd4j.getExecutioner().exec(op);
+
+        norm.addi(epsilon);
+        clipped.divi(norm);
+
+        if(rate != 1.0){
+            clipped.muli(rate).addi(norm.muli(1.0-rate));
+        }
+
+        Broadcast.mul(param, clipped, param, getBroadcastDims(dimensions, param.rank()) );
+    }
+
+    @Override
+    public MinMaxNormConstraint clone() {
+        return new MinMaxNormConstraint(min, max, rate, params, dimensions);
+    }
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/conf/constraint/NonNegativeConstraint.java

@@ -0,0 +1,43 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.nn.conf.constraint;
+
+import lombok.Data;
+import lombok.EqualsAndHashCode;
+import org.nd4j.linalg.api.ndarray.INDArray;
+import org.nd4j.linalg.indexing.BooleanIndexing;
+import org.nd4j.linalg.indexing.conditions.Conditions;
+
+@Data
+@EqualsAndHashCode(callSuper = true)
+public class NonNegativeConstraint extends BaseConstraint {
+
+    public NonNegativeConstraint(){ }
+
+    @Override
+    public void apply(INDArray param) {
+        BooleanIndexing.replaceWhere(param, 0.0, Conditions.lessThan(0.0));
+    }
+
+    @Override
+    public NonNegativeConstraint clone() { return new NonNegativeConstraint();}
+
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/conf/constraint/UnitNormConstraint.java

@@ -0,0 +1,70 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.nn.conf.constraint;
+
+import lombok.Data;
+import lombok.EqualsAndHashCode;
+import org.nd4j.linalg.api.ndarray.INDArray;
+import org.nd4j.linalg.factory.Broadcast;
+
+import java.util.Collections;
+import java.util.Set;
+
+@Data
+@EqualsAndHashCode(callSuper = true)
+public class UnitNormConstraint extends BaseConstraint {
+
+    private UnitNormConstraint(){
+        //No arg for json ser/de
+    }
+
+    /**
+     * Apply to weights but not biases by default
+     *
+     * @param dimensions     Dimensions to apply to. For DenseLayer, OutputLayer, RnnOutputLayer, LSTM, etc: this should
+     *                       be dimension 1. For CNNs, this should be dimensions [1,2,3] corresponding to last 3 of
+     *                       parameters which have order [depthOut, depthIn, kH, kW]
+     */
+    public UnitNormConstraint(long... dimensions){
+        this(Collections.<String>emptySet(), dimensions);
+    }
+
+
+    /**
+     * @param dimensions     Dimensions to apply to. For DenseLayer, OutputLayer, RnnOutputLayer, LSTM, etc: this should
+     *                       be dimension 1. For CNNs, this should be dimensions [1,2,3] corresponding to last 3 of
+     *                       parameters which have order [depthOut, depthIn, kH, kW]
+     */
+    public UnitNormConstraint(Set<String> paramNames, long... dimensions){
+        super(paramNames, dimensions);
+    }
+
+    @Override
+    public void apply(INDArray param) {
+        INDArray norm2 = param.norm2(dimensions);
+        Broadcast.div(param, norm2, param, getBroadcastDims(dimensions, param.rank()) );
+    }
+
+    @Override
+    public UnitNormConstraint clone() {
+        return new UnitNormConstraint( params, dimensions);
+    }
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/conf/distribution/BinomialDistribution.java

@@ -0,0 +1,89 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.nn.conf.distribution;
+
+import org.nd4j.shade.jackson.annotation.JsonCreator;
+import org.nd4j.shade.jackson.annotation.JsonProperty;
+
+public class BinomialDistribution extends Distribution {
+
+    private static final long serialVersionUID = 7407024251874318749L;
+
+    private final int numberOfTrials;
+    private double probabilityOfSuccess;
+
+    /**
+     * Create a distribution
+     *
+     * @param numberOfTrials the number of trials
+     * @param probabilityOfSuccess the probability of success
+     */
+    @JsonCreator
+    public BinomialDistribution(@JsonProperty("numberOfTrials") int numberOfTrials,
+                    @JsonProperty("probabilityOfSuccess") double probabilityOfSuccess) {
+        this.numberOfTrials = numberOfTrials;
+        this.probabilityOfSuccess = probabilityOfSuccess;
+    }
+
+    public double getProbabilityOfSuccess() {
+        return probabilityOfSuccess;
+    }
+
+    public void setProbabilityOfSuccess(double probabilityOfSuccess) {
+        this.probabilityOfSuccess = probabilityOfSuccess;
+    }
+
+    public int getNumberOfTrials() {
+        return numberOfTrials;
+    }
+
+    @Override
+    public int hashCode() {
+        final int prime = 31;
+        int result = 1;
+        result = prime * result + numberOfTrials;
+        long temp;
+        temp = Double.doubleToLongBits(probabilityOfSuccess);
+        result = prime * result + (int) (temp ^ (temp >>> 32));
+        return result;
+    }
+
+    @Override
+    public boolean equals(Object obj) {
+        if (this == obj)
+            return true;
+        if (obj == null)
+            return false;
+        if (getClass() != obj.getClass())
+            return false;
+        BinomialDistribution other = (BinomialDistribution) obj;
+        if (numberOfTrials != other.numberOfTrials)
+            return false;
+        if (Double.doubleToLongBits(probabilityOfSuccess) != Double.doubleToLongBits(other.probabilityOfSuccess))
+            return false;
+        return true;
+    }
+
+    public String toString() {
+        return "BinomialDistribution(" + "numberOfTrials=" + numberOfTrials + ", probabilityOfSuccess="
+                        + probabilityOfSuccess + ')';
+    }
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/conf/distribution/ConstantDistribution.java

@@ -0,0 +1,47 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.nn.conf.distribution;
+
+import lombok.Data;
+import lombok.EqualsAndHashCode;
+import org.nd4j.shade.jackson.annotation.JsonCreator;
+import org.nd4j.shade.jackson.annotation.JsonProperty;
+
+@Data
+@EqualsAndHashCode(callSuper = false)
+public class ConstantDistribution extends Distribution {
+
+    private double value;
+
+    /**
+     * Create a Constant distribution with given value
+     *
+     * @param value the gain
+     */
+    @JsonCreator
+    public ConstantDistribution(@JsonProperty("value") double value) {
+        this.value = value;
+    }
+
+    public String toString() {
+        return "ConstantDistribution(value=" + value + ")";
+    }
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/conf/distribution/Distribution.java

@@ -0,0 +1,42 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.nn.conf.distribution;
+
+import org.deeplearning4j.nn.conf.distribution.serde.LegacyDistributionHelper;
+import org.nd4j.shade.jackson.annotation.JsonTypeInfo;
+
+import java.io.Serializable;
+
+@JsonTypeInfo(use = JsonTypeInfo.Id.CLASS, include = JsonTypeInfo.As.PROPERTY, property = "type",
+                defaultImpl = LegacyDistributionHelper.class)
+public abstract class Distribution implements Serializable, Cloneable {
+
+    private static final long serialVersionUID = 5401741214954998498L;
+
+    @Override
+    public Distribution clone() {
+        try {
+            return (Distribution) super.clone();
+        } catch (CloneNotSupportedException e) {
+            throw new RuntimeException(e);
+        }
+    }
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/conf/distribution/Distributions.java

@@ -0,0 +1,65 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.nn.conf.distribution;
+
+import org.nd4j.linalg.factory.Nd4j;
+
+public class Distributions {
+    private Distributions() {}
+
+    public static org.nd4j.linalg.api.rng.distribution.Distribution createDistribution(Distribution dist) {
+        if (dist == null)
+            return null;
+        if (dist instanceof NormalDistribution) {
+            NormalDistribution nd = (NormalDistribution) dist;
+            return Nd4j.getDistributions().createNormal(nd.getMean(), nd.getStd());
+        }
+        if (dist instanceof GaussianDistribution) {
+            GaussianDistribution nd = (GaussianDistribution) dist;
+            return Nd4j.getDistributions().createNormal(nd.getMean(), nd.getStd());
+        }
+        if (dist instanceof UniformDistribution) {
+            UniformDistribution ud = (UniformDistribution) dist;
+            return Nd4j.getDistributions().createUniform(ud.getLower(), ud.getUpper());
+        }
+        if (dist instanceof BinomialDistribution) {
+            BinomialDistribution bd = (BinomialDistribution) dist;
+            return Nd4j.getDistributions().createBinomial(bd.getNumberOfTrials(), bd.getProbabilityOfSuccess());
+        }
+        if (dist instanceof LogNormalDistribution) {
+            LogNormalDistribution lnd = (LogNormalDistribution) dist;
+            return Nd4j.getDistributions().createLogNormal(lnd.getMean(), lnd.getStd());
+        }
+        if (dist instanceof TruncatedNormalDistribution) {
+            TruncatedNormalDistribution tnd = (TruncatedNormalDistribution) dist;
+            return Nd4j.getDistributions().createTruncatedNormal(tnd.getMean(), tnd.getStd());
+        }
+        if (dist instanceof OrthogonalDistribution) {
+            OrthogonalDistribution od = (OrthogonalDistribution) dist;
+            return Nd4j.getDistributions().createOrthogonal(od.getGain());
+        }
+        if (dist instanceof ConstantDistribution) {
+            ConstantDistribution od = (ConstantDistribution) dist;
+            return Nd4j.getDistributions().createConstant(od.getValue());
+        }
+        throw new RuntimeException("unknown distribution type: " + dist.getClass());
+    }
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/conf/distribution/GaussianDistribution.java

@@ -0,0 +1,40 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.nn.conf.distribution;
+
+import org.nd4j.shade.jackson.annotation.JsonCreator;
+import org.nd4j.shade.jackson.annotation.JsonProperty;
+
+@Deprecated
+public class GaussianDistribution extends NormalDistribution {
+
+    /**
+     * Create a gaussian distribution (equivalent to normal)
+     * with the given mean and std
+     *
+     * @param mean the mean
+     * @param std  the standard deviation
+     */
+    @JsonCreator
+    public GaussianDistribution(@JsonProperty("mean") double mean, @JsonProperty("std") double std) {
+        super(mean, std);
+    }
+}

deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/conf/distribution/LogNormalDistribution.java

@@ -0,0 +1,56 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.nn.conf.distribution;
+
+import lombok.Data;
+import lombok.EqualsAndHashCode;
+import org.nd4j.shade.jackson.annotation.JsonCreator;
+import org.nd4j.shade.jackson.annotation.JsonProperty;
+
+/**
+ * A log-normal distribution, with two parameters: mean and standard deviation.
+ * Note: the mean and standard deviation are for the logarithm of the values.
+ * Put another way: if X~LogN(M,S), then mean(log(X))=M, and stdev(log(X))=S
+ *
+ */
+@EqualsAndHashCode(callSuper = false)
+@Data
+public class LogNormalDistribution extends Distribution {
+
+    private double mean, std;
+
+    /**
+     * Create a log-normal distribution
+     * with the given mean and std
+     *
+     * @param mean the mean
+     * @param std  the standard deviation
+     */
+    @JsonCreator
+    public LogNormalDistribution(@JsonProperty("mean") double mean, @JsonProperty("std") double std) {
+        this.mean = mean;
+        this.std = std;
+    }
+
+    public String toString() {
+        return "LogNormalDistribution(" + "mean=" + mean + ", std=" + std + ')';
+    }
+}