chore: import upstream snapshot with attribution

deeplearning4j/deeplearning4j-parallelwrapper/pom.xml

@@ -0,0 +1,85 @@
+<?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-parallel-wrapper</artifactId>
+
+    <name>deeplearning4j-parallel-wrapper</name>
+
+    <build>
+        <plugins>
+            <plugin>
+                <groupId>org.moditect</groupId>
+                <artifactId>moditect-maven-plugin</artifactId>
+            </plugin>
+        </plugins>
+    </build>
+
+    <properties>
+
+
+        <module.name>deeplearning4j.parallel.wrapper</module.name>
+    </properties>
+
+    <dependencies>
+        <dependency>
+            <groupId>com.beust</groupId>
+            <artifactId>jcommander</artifactId>
+            <version>${jcommander.version}</version>
+        </dependency>
+        <!-- Logging Dependencies -->
+        <dependency>
+            <groupId>org.slf4j</groupId>
+            <artifactId>slf4j-api</artifactId>
+        </dependency>
+        <!-- Redirect jackson to slf4j.  -->
+        <dependency>
+            <groupId>org.slf4j</groupId>
+            <artifactId>log4j-over-slf4j</artifactId>
+        </dependency>
+
+        <dependency>
+            <groupId>org.eclipse.deeplearning4j</groupId>
+            <artifactId>deeplearning4j-core</artifactId>
+            <version>${project.version}</version>
+        </dependency>
+        <dependency>
+            <groupId>org.eclipse.deeplearning4j</groupId>
+            <artifactId>deeplearning4j-ui</artifactId>
+            <version>${project.version}</version>
+            <scope>test</scope>
+        </dependency>
+
+    </dependencies>
+
+
+</project>

deeplearning4j/deeplearning4j-parallelwrapper/src/main/java/org/deeplearning4j/parallelism/InplaceParallelInference.java

@@ -0,0 +1,330 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.parallelism;
+
+
+import lombok.AllArgsConstructor;
+import lombok.NoArgsConstructor;
+import lombok.NonNull;
+import lombok.extern.slf4j.Slf4j;
+import lombok.val;
+
+import org.deeplearning4j.nn.api.Model;
+import org.deeplearning4j.nn.api.ModelAdapter;
+import org.deeplearning4j.nn.conf.ComputationGraphConfiguration;
+import org.deeplearning4j.nn.conf.MultiLayerConfiguration;
+import org.deeplearning4j.nn.graph.ComputationGraph;
+import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;
+import org.deeplearning4j.parallelism.inference.InferenceMode;
+import org.deeplearning4j.parallelism.inference.LoadBalanceMode;
+import org.nd4j.linalg.api.ndarray.INDArray;
+import org.nd4j.linalg.exception.ND4JIllegalStateException;
+import org.nd4j.linalg.factory.Nd4j;
+
+import java.util.*;
+import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.CopyOnWriteArrayList;
+import java.util.concurrent.LinkedBlockingQueue;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
+
+@Slf4j
+public class InplaceParallelInference extends ParallelInference {
+    protected List<ModelHolder> holders = new CopyOnWriteArrayList<>();
+    protected String[] layersToOutputTo;
+    protected int[] layerIndicesOutputTo;
+
+    protected ModelSelector selector = new ModelSelector();
+
+    protected final Object locker = new Object();
+
+    @Override
+    protected void init() {
+        for (int e = 0; e < Nd4j.getAffinityManager().getNumberOfDevices(); e++) {
+            val h = ModelHolder.builder()
+                    .sourceModel(model)
+                    .workers(workers)
+                    .layerIndicesOutputTo(layerIndicesOutputTo)
+                    .layersToOutputTo(layersToOutputTo)
+                    .loadBalanceMode(loadBalanceMode)
+                    .targetDeviceId(e)
+                    .rootDevice(e == Nd4j.getAffinityManager().getDeviceForCurrentThread().intValue())
+                    .build();
+            h.init();
+
+            // adding for simplified access
+            holders.add(h);
+
+            // and adding it to actual
+            selector.addModelHolder(e, h);
+        }
+    }
+
+    @Override
+    public synchronized void updateModel(@NonNull Model model) {
+        for (val h:holders)
+            h.updateModel(model);
+    }
+
+    @Override
+    protected synchronized Model[] getCurrentModelsFromWorkers() {
+        val models = new Model[holders.size()];
+        int cnt = 0;
+        for (val h:holders) {
+            models[cnt++] = h.sourceModel;
+        }
+
+        return models;
+    }
+
+    @Override
+    public INDArray[] output(INDArray[] input, INDArray[] inputMasks) {
+        return selector.output(input, inputMasks);
+    }
+
+    /**
+     * This method does forward pass and returns output provided by OutputAdapter
+     *
+     * @param adapter
+     * @param input
+     * @param inputMasks
+     * @param <T>
+     * @return
+     */
+    public <T> T output(@NonNull ModelAdapter<T> adapter, INDArray[] input, INDArray[] inputMasks, INDArray[] labelsMasks) {
+        val holder = selector.getModelForThisThread();
+        Model model = null;
+        boolean acquired = false;
+        try {
+            model = holder.acquireModel();
+            acquired = true;
+            return adapter.apply(model, input, inputMasks, labelsMasks);
+        } catch (InterruptedException e) {
+            throw new RuntimeException(e);
+        } finally {
+            if (model != null && acquired)
+                holder.releaseModel(model);
+        }
+    }
+
+
+    protected static class ModelSelector {
+        // this map stores collection of shared
+        protected Map<Integer, ModelHolder> map = new HashMap<>();
+
+        protected final LoadBalanceMode loadBalanceMode;
+
+        public ModelSelector() {
+            this(LoadBalanceMode.ROUND_ROBIN);
+        }
+
+        public ModelSelector(LoadBalanceMode loadBalanceMode) {
+            this.loadBalanceMode = loadBalanceMode;
+        }
+
+        protected void addModelHolder(@NonNull Integer device, @NonNull ModelHolder holder) {
+            map.put(device, holder);
+        }
+
+        public ModelHolder getModelForThread(long threadId) {
+            // first of all we get mapped device for this thread
+            val device = Nd4j.getAffinityManager().getDeviceForThread(threadId);
+
+            // each device has it's own queue
+            val q = map.get(device);
+
+            // and we're returning holder right away
+            return q;
+        }
+
+        public INDArray[] output(INDArray[] input, INDArray[] inputMasks) {
+            return getModelForThisThread().output(input, inputMasks);
+        }
+
+        public ModelHolder getModelForThisThread() {
+            return getModelForThread(Thread.currentThread().getId());
+        }
+    }
+
+    @NoArgsConstructor
+    @AllArgsConstructor
+    @lombok.Builder
+    protected static class ModelHolder {
+        protected Model sourceModel;
+        @lombok.Builder.Default protected int workers = 4;
+        @lombok.Builder.Default protected List<Model> replicas = new ArrayList<>();
+        @lombok.Builder.Default protected boolean rootDevice = true;
+        @lombok.Builder.Default protected LoadBalanceMode loadBalanceMode = LoadBalanceMode.ROUND_ROBIN;
+
+        protected String[] layersToOutputTo;
+        protected int[] layerIndicesOutputTo;
+        protected int targetDeviceId;
+
+        protected final AtomicLong position = new AtomicLong(0);
+        protected final ReentrantReadWriteLock modelLock = new ReentrantReadWriteLock();
+
+        // this queue is used in FIFO mode
+        protected final BlockingQueue<Model> queue = new LinkedBlockingQueue<>();
+
+        @lombok.Builder.Default protected transient boolean isCG = false;
+        @lombok.Builder.Default protected transient boolean isMLN = false;
+
+
+        protected synchronized void init() {
+            if (workers < 1)
+                throw new ND4JIllegalStateException("Workers must be positive value");
+
+            replicas.clear();
+
+            isCG = sourceModel instanceof ComputationGraph;
+            isMLN = sourceModel instanceof MultiLayerNetwork;
+
+            // we clone params only if we're not on the same device
+            val params = rootDevice ? sourceModel.params() : sourceModel.params().unsafeDuplication(true);
+
+            // and moving it to specified device (only if NOT root
+            if (!rootDevice)
+                Nd4j.getAffinityManager().replicateToDevice(targetDeviceId, params);
+
+            for (int e = 0; e < workers; e++) {
+                if (sourceModel instanceof ComputationGraph) {
+                    // building configuration with shared parameters
+                    val model = new ComputationGraph(ComputationGraphConfiguration.fromJson(((ComputationGraph) sourceModel).getConfiguration().toJson()));
+                    model.init(params, false);
+                    Nd4j.getExecutioner().commit();
+
+                    // storing model for future reuse
+                    replicas.add(model);
+
+                    if (loadBalanceMode == LoadBalanceMode.FIFO)
+                        queue.add(model);
+                } else if (sourceModel instanceof MultiLayerNetwork) {
+                    val model = new MultiLayerNetwork(MultiLayerConfiguration.fromJson(((MultiLayerNetwork) sourceModel).getLayerWiseConfigurations().toJson()));
+                    model.init(params, false);
+                    Nd4j.getExecutioner().commit();
+
+                    replicas.add(model);
+
+                    if (loadBalanceMode == LoadBalanceMode.FIFO)
+                        queue.add(model);
+                }
+            }
+        }
+
+
+        protected Model acquireModel() throws InterruptedException {
+            try {
+                modelLock.readLock().lock();
+
+                switch (loadBalanceMode) {
+                    case FIFO: {
+                        return queue.take();
+                    }
+                    case ROUND_ROBIN:
+                        return replicas.get((int) (position.getAndIncrement() % replicas.size()));
+                    default:
+                        throw new ND4JIllegalStateException("Unknown LoadBalanceMode was specified: [" + loadBalanceMode + "]");
+                }
+            } finally {
+                modelLock.readLock().unlock();
+            }
+        }
+
+        protected void releaseModel(Model model) {
+            try {
+                modelLock.readLock().lock();
+
+                switch (loadBalanceMode) {
+                    case FIFO:
+                        queue.add(model);
+                        break;
+                    case ROUND_ROBIN:
+                        break;
+                    default:
+                        throw new ND4JIllegalStateException("Unknown LoadBalanceMode was specified: [" + loadBalanceMode + "]");
+                }
+            } finally {
+                modelLock.readLock().unlock();
+            }
+        }
+
+        protected INDArray[] output(INDArray[] input, INDArray[] inputMasks) {
+            try {
+                modelLock.readLock().lock();
+                if (isCG) {
+                    // acquiring model from pool
+                    val model = acquireModel();
+
+                    // doing inference
+                    INDArray[] output;
+                    try{
+                        if(layersToOutputTo != null) {
+                            output = ((ComputationGraph) model).output(Arrays.asList(layersToOutputTo),false,input,inputMasks);
+                        }
+                        else
+                            output = ((ComputationGraph) model).output(false, input, inputMasks);
+                    } finally {
+                        // releasing model
+                        releaseModel(model);
+                    }
+                    return output;
+                } else if (isMLN) {
+                    if (input.length > 1 || (inputMasks != null && inputMasks.length > 1))
+                        throw new ND4JIllegalStateException("MultilayerNetwork can't have multiple inputs");
+
+                    val model = acquireModel();
+                    INDArray result;
+                    try {
+                        if(layerIndicesOutputTo != null) {
+                            MultiLayerNetwork multiLayerNetwork = (MultiLayerNetwork) model;
+                            result = multiLayerNetwork.feedForwardToLayer(layerIndicesOutputTo[0],input[0],false).get(0);
+
+                        } else {
+                            result = ((MultiLayerNetwork) model).output(input[0], false, (inputMasks == null ? null : inputMasks[0]), null);
+
+                        }
+                    } finally {
+                        releaseModel(model);
+                    }
+                    return new INDArray[]{result};
+                } else
+                    throw new UnsupportedOperationException();
+            } catch (InterruptedException e) {
+                throw new RuntimeException(e);
+            } finally {
+                modelLock.readLock().unlock();
+            }
+        }
+
+        protected void updateModel(@NonNull Model model) {
+            try {
+                modelLock.writeLock().lock();
+
+                this.sourceModel = model;
+
+                init();
+            } finally {
+                modelLock.writeLock().unlock();
+            }
+        }
+
+    }
+}

deeplearning4j/deeplearning4j-parallelwrapper/src/main/java/org/deeplearning4j/parallelism/ParallelInference.java

@@ -0,0 +1,690 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.parallelism;
+
+import lombok.NonNull;
+import lombok.extern.slf4j.Slf4j;
+import org.deeplearning4j.nn.api.Model;
+import org.deeplearning4j.nn.api.ModelAdapter;
+import org.deeplearning4j.nn.conf.ComputationGraphConfiguration;
+import org.deeplearning4j.nn.conf.MultiLayerConfiguration;
+import org.deeplearning4j.nn.graph.ComputationGraph;
+import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;
+import org.deeplearning4j.parallelism.inference.InferenceMode;
+import org.deeplearning4j.parallelism.inference.InferenceObservable;
+import org.deeplearning4j.parallelism.inference.LoadBalanceMode;
+import org.deeplearning4j.parallelism.inference.observers.BasicInferenceObservable;
+import org.deeplearning4j.parallelism.inference.observers.BasicInferenceObserver;
+import org.deeplearning4j.parallelism.inference.observers.BatchedInferenceObservable;
+import org.nd4j.linalg.api.ndarray.INDArray;
+import org.nd4j.linalg.dataset.DataSet;
+import org.nd4j.linalg.exception.ND4JIllegalStateException;
+import org.nd4j.linalg.factory.Nd4j;
+import org.nd4j.common.primitives.Pair;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+import java.util.Observer;
+import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.LinkedBlockingQueue;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
+
+@Slf4j
+public class ParallelInference {
+    protected Model model;
+    protected long nanos;
+    protected int workers;
+    protected int batchLimit;
+    protected InferenceMode inferenceMode;
+    protected int queueLimit;
+    protected LoadBalanceMode loadBalanceMode = LoadBalanceMode.FIFO;
+
+    // this queue holds data for inference
+    private BlockingQueue<InferenceObservable> observables;
+
+    private final Object locker = new Object();
+
+    private InferenceWorker[] zoo;
+    private ObservablesProvider provider;
+
+    protected String[] layersToOutputTo;
+    protected int[] layerIndicesOutputTo;
+
+
+    public final static int DEFAULT_NUM_WORKERS = Nd4j.getAffinityManager().getNumberOfDevices();
+    public final static int DEFAULT_BATCH_LIMIT = 32;
+    public final static InferenceMode DEFAULT_INFERENCE_MODE = InferenceMode.BATCHED;
+    public final static int DEFAULT_QUEUE_LIMIT = 64;
+
+
+
+    protected ParallelInference() {
+        //
+    }
+
+    /**
+     * This method allows to update Model used for inference in runtime, without queue reset
+     *
+     * @param model
+     */
+    public void updateModel(@NonNull Model model) {
+        if (zoo != null) {
+            for (var w: zoo)
+                w.updateModel(model);
+        } else {
+            // if zoo wasn't initialized yet - just replace model
+            this.model = model;
+        }
+    }
+
+    /**
+     * This method returns Models used in workers at this moment
+     * PLEASE NOTE: This method is NOT thread safe, and should NOT be used anywhere but tests
+     *
+     * @return
+     */
+    protected Model[] getCurrentModelsFromWorkers() {
+        if (zoo == null)
+            return new Model[0];
+
+        var models = new Model[zoo.length];
+        int cnt = 0;
+        for (var w:zoo) {
+            models[cnt++] = w.replicatedModel;
+        }
+
+        return models;
+    }
+
+    protected void init() {
+        observables = new LinkedBlockingQueue<>(queueLimit);
+
+        int numDevices = Nd4j.getAffinityManager().getNumberOfDevices();
+        int currentDevice = Nd4j.getAffinityManager().getDeviceForCurrentThread();
+        AtomicBoolean assignedRoot = new AtomicBoolean(false);
+
+        zoo = new InferenceWorker[workers];
+        for (int i = 0; i < workers; i++) {
+            int cDevice = i % numDevices;
+            boolean cRoot = !assignedRoot.get() && cDevice == currentDevice;
+            assignedRoot.compareAndSet(false, cRoot);
+            if(layersToOutputTo != null)
+                zoo[i] = new InferenceWorker(layersToOutputTo,i, model, observables, cRoot, cDevice);
+            else if(layerIndicesOutputTo != null)
+                zoo[i] = new InferenceWorker(layerIndicesOutputTo,i, model, observables, cRoot, cDevice);
+            else
+                zoo[i] = new InferenceWorker(i, model, observables, cRoot, cDevice);
+
+
+            zoo[i].setDaemon(true);
+            zoo[i].start();
+        }
+
+
+        if (inferenceMode == InferenceMode.BATCHED) {
+            log.info("Initializing ObservablesProvider...");
+            provider = new ObservablesProvider(nanos, batchLimit, observables);
+        }
+    }
+
+    protected long getWorkerCounter(int workerIdx) {
+        return zoo[workerIdx].getCounterValue();
+    }
+
+    /**
+     * This method gracefully shuts down ParallelInference instance
+     */
+    public synchronized void shutdown() {
+        if (zoo == null)
+            return;
+
+        for (int e = 0; e < zoo.length; e++) {
+            if (zoo[e] == null)
+                continue;
+
+            zoo[e].interrupt();
+            zoo[e].shutdown();
+            zoo[e] = null;
+        }
+        zoo = null;
+
+        System.gc();
+    }
+
+    /**
+     *
+     * @param input
+     * @return
+     */
+    public INDArray output(double[] input) {
+        return output(Nd4j.create(input));
+    }
+
+    /**
+     *
+     * @param input
+     * @return
+     */
+    public INDArray output(float[] input) {
+        return output(Nd4j.create(input));
+    }
+
+    public INDArray output(INDArray input) {
+        return output(input, null);
+    }
+
+    public INDArray output(INDArray input, INDArray inputMask) {
+        INDArray[] out = output(new INDArray[]{input}, (inputMask == null ? null : new INDArray[]{inputMask}));
+        // basically, depending on model type we either
+        // throw stuff to specific model, or wait for batch
+        if(out.length != 1){
+            throw new IllegalArgumentException("Network has multiple (" + out.length + ") output arrays, but only a" +
+                    " single output can be returned using this method. Use for output(INDArray[] input, INDArray[] " +
+                    "inputMasks) for multi-output nets");
+        }
+        return out[0];
+    }
+
+    /**
+     *
+     * @param dataSet
+     * @return
+     */
+    public INDArray output(DataSet dataSet) {
+        return output(dataSet.getFeatures(), dataSet.getFeaturesMaskArray());
+    }
+
+    /**
+     * Generate predictions/output from the network
+     *
+     * @param input Input to the network
+     * @return Output from the network
+     */
+    public INDArray[] output(INDArray... input) {
+        return output(input, null);
+    }
+
+    /**
+     * Generate predictions/outputs from the network, optionally using input masks for predictions
+     *
+     * @param input      Input to the network
+     * @param inputMasks Input masks for the network. May be null.
+     * @return Output from the network
+     */
+    public INDArray[] output(INDArray[] input, INDArray[] inputMasks){
+        Nd4j.getExecutioner().commit(); //Commit before passing input to other thread
+
+        // basically, depending on model type we either throw stuff to specific model, or wait for batch
+        BasicInferenceObserver observer = new BasicInferenceObserver();
+        InferenceObservable observable;
+
+        if (inferenceMode == InferenceMode.SEQUENTIAL) {
+            if(layersToOutputTo != null)
+                observable = new BasicInferenceObservable(layersToOutputTo,input, inputMasks);
+            else if(layerIndicesOutputTo != null)
+                observable = new BasicInferenceObservable(layerIndicesOutputTo,input, inputMasks);
+            else
+                observable = new BasicInferenceObservable(input, inputMasks);
+
+            observable.addObserver(observer);
+            try {
+                observables.put(observable);
+            } catch (InterruptedException e) {
+                Thread.currentThread().interrupt();
+                throw new RuntimeException(e);
+            }
+        } else {
+            observable = provider.setInput(observer, input, inputMasks);
+        }
+
+        try {
+            // submit query to processing
+            // and block until Observable returns
+            //observer.wait();
+
+            observer.waitTillDone();
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+
+        return observable.getOutput();
+    }
+
+    /**
+     * This method does forward pass and returns output provided by OutputAdapter
+     *
+     * @param adapter
+     * @param inputs
+     * @return
+     */
+    public <T> T output(@NonNull ModelAdapter<T> adapter, INDArray... inputs) {
+        return output(adapter, inputs, null);
+    }
+
+    /**
+     * This method does forward pass and returns output provided by OutputAdapter
+     *
+     * @param adapter
+     * @param input
+     * @param inputMasks
+     * @param <T>
+     * @return
+     */
+    public <T> T output(@NonNull ModelAdapter<T> adapter,INDArray[] input, INDArray[] inputMasks) {
+        throw new ND4JIllegalStateException("Adapted mode requires Inplace inference mode");
+    }
+
+
+    public static class Builder {
+        private Model model;
+        private int workers = DEFAULT_NUM_WORKERS;
+        private int batchLimit = DEFAULT_BATCH_LIMIT;
+        private InferenceMode inferenceMode = DEFAULT_INFERENCE_MODE;
+        private int queueLimit = DEFAULT_QUEUE_LIMIT;
+        private String[] layersToOutputTo;
+        private int[] layerIndicesOutputTo;
+
+        protected LoadBalanceMode loadBalanceMode = LoadBalanceMode.FIFO;
+
+        public Builder(@NonNull Model model) {
+            this.model = model;
+        }
+
+        /**
+         * Optional input for outputting to a subset of layers according to indices.
+         * Used in {@link MultiLayerNetwork}
+         * @param layerIndicesOutputTo the layer indices
+         * @return
+         */
+        public Builder layerIndicesOutputTo(int[] layerIndicesOutputTo) {
+            this.layerIndicesOutputTo = layerIndicesOutputTo;
+            return this;
+        }
+
+        /**
+         * Optional input for outputting to a subset of layers according to indices.
+         * Used in {@link ComputationGraph}
+         * @param layersToOutputTo the layer output names
+         * @return
+         */
+        public Builder layersToOutputTo(String[] layersToOutputTo) {
+            this.layersToOutputTo = layersToOutputTo;
+            return this;
+        }
+
+        /**
+         * This method allows you to define mode that'll be used during inference. Options are:
+         *
+         * SEQUENTIAL: Input will be sent to last-used worker unmodified.
+         * BATCHED: Multiple inputs will be packed into single batch, and
+         * sent to last-used device.
+         *
+         * @param inferenceMode
+         * @return
+         */
+        public Builder inferenceMode(@NonNull InferenceMode inferenceMode) {
+            this.inferenceMode = inferenceMode;
+            return this;
+        }
+
+
+        /**
+         * This method allows you to specify load balance mode
+         *
+         * @param loadBalanceMode
+         * @return
+         */
+        public Builder loadBalanceMode(@NonNull LoadBalanceMode loadBalanceMode) {
+            this.loadBalanceMode = loadBalanceMode;
+            return this;
+        }
+
+
+        /**
+         * This method defines, how many model copies will be used for inference.
+         *
+         * PLEASE NOTE: This method primarily suited for multi-GPU systems
+         * PLEASE NOTE: For INPLACE inference mode this value will mean number of models per DEVICE
+         *
+         * @param workers
+         * @return
+         */
+        public Builder workers(int workers) {
+            if (workers < 1)
+                throw new IllegalStateException("Workers should be positive value");
+
+            this.workers = workers;
+            return this;
+        }
+
+        /**
+         * This method defines, how many input samples can
+         * be batched within given time frame.
+         *
+         * PLEASE NOTE: This value has no effect in
+         * SEQUENTIAL inference mode
+         *
+         * @param limit
+         * @return
+         */
+        public Builder batchLimit(int limit) {
+            if (limit < 1)
+                throw new IllegalStateException("Batch limit should be positive value");
+
+            this.batchLimit = limit;
+            return this;
+        }
+
+        /**
+         * This method defines buffer queue size.
+         *
+         * Default value: 64
+         *
+         * @param limit
+         * @return
+         */
+        public Builder queueLimit(int limit) {
+            if (limit < 1)
+                throw new IllegalStateException("Queue limit should be positive value");
+
+            this.queueLimit = limit;
+            return this;
+        }
+
+        /**
+         * This method builds new ParallelInference instance
+         *
+         * @return
+         */
+        public ParallelInference build() {
+            if (this.inferenceMode == InferenceMode.INPLACE) {
+                var inf = new InplaceParallelInference();
+                inf.inferenceMode = this.inferenceMode;
+                inf.model = this.model;
+                inf.layersToOutputTo = this.layersToOutputTo;
+                inf.layerIndicesOutputTo = this.layerIndicesOutputTo;
+                inf.workers = this.workers;
+                inf.loadBalanceMode = this.loadBalanceMode;
+                inf.init();
+
+                return inf;
+            } else {
+                ParallelInference inference = new ParallelInference();
+                inference.batchLimit = this.batchLimit;
+                inference.queueLimit = this.queueLimit;
+                inference.inferenceMode = this.inferenceMode;
+                inference.model = this.model;
+                inference.workers = this.workers;
+                inference.loadBalanceMode = this.loadBalanceMode;
+                inference.layerIndicesOutputTo = layerIndicesOutputTo;
+                inference.layersToOutputTo = layersToOutputTo;
+                inference.init();
+
+                return inference;
+            }
+        }
+    }
+
+
+    /**
+     * This class actually does inference with respect to device affinity
+     *
+     */
+    private class InferenceWorker extends Thread implements Runnable {
+        private BlockingQueue<InferenceObservable> inputQueue;
+        private AtomicBoolean shouldWork = new AtomicBoolean(true);
+        private AtomicBoolean isStopped = new AtomicBoolean(false);
+        private Model protoModel;
+        private Model replicatedModel;
+        private AtomicLong counter = new AtomicLong(0);
+        private boolean rootDevice;
+        private int deviceId;
+
+        private String[] layersToOutputTo;
+        private int[] layerIndicesOutputTo;
+
+        private ReentrantReadWriteLock modelLock = new ReentrantReadWriteLock();
+
+        private InferenceWorker(String[] layersToOutputTo,int id, @NonNull Model model, @NonNull BlockingQueue inputQueue, boolean rootDevice, int deviceId) {
+            this.inputQueue = inputQueue;
+            this.protoModel = model;
+            this.rootDevice = rootDevice;
+            this.deviceId = deviceId;
+            this.layersToOutputTo = layersToOutputTo;
+            this.setDaemon(true);
+            this.setName("InferenceThread-" + id);
+
+        }
+
+
+        private InferenceWorker(int[] layerIndicesOutputTo,int id, @NonNull Model model, @NonNull BlockingQueue inputQueue, boolean rootDevice, int deviceId) {
+            this(id,model,inputQueue,rootDevice,deviceId);
+            this.layerIndicesOutputTo = layerIndicesOutputTo;
+
+        }
+
+        private InferenceWorker(int id, @NonNull Model model, @NonNull BlockingQueue inputQueue, boolean rootDevice, int deviceId) {
+            this.inputQueue = inputQueue;
+            this.protoModel = model;
+            this.rootDevice = rootDevice;
+            this.deviceId = deviceId;
+
+            this.setDaemon(true);
+            this.setName("InferenceThread-" + id);
+
+        }
+        protected long getCounterValue() {
+            return counter.get();
+        }
+
+        protected void updateModel(@NonNull Model model) {
+            try {
+                modelLock.writeLock().lock();
+                this.protoModel = model;
+
+                // now re-init model
+                initializeReplicaModel();
+            } finally {
+                modelLock.writeLock().unlock();
+            }
+        }
+
+        /**
+         * This method duplicates model for future use during inference
+         */
+        protected void initializeReplicaModel() {
+            if (protoModel instanceof ComputationGraph) {
+                if (!rootDevice) {
+                    this.replicatedModel = new ComputationGraph(ComputationGraphConfiguration
+                            .fromJson(((ComputationGraph) protoModel).getConfiguration().toJson()));
+                    this.replicatedModel.init();
+
+                    synchronized (locker) {
+                        this.replicatedModel.setParams(protoModel.params().unsafeDuplication(true));
+
+                        Nd4j.getExecutioner().commit();
+                    }
+                } else {
+                    this.replicatedModel = protoModel;
+                }
+            } else if (protoModel instanceof MultiLayerNetwork) {
+                if (!rootDevice) {
+                    this.replicatedModel = new MultiLayerNetwork(MultiLayerConfiguration.fromJson(
+                            ((MultiLayerNetwork) protoModel).getLayerWiseConfigurations().toJson()));
+                    this.replicatedModel.init();
+
+                    synchronized (locker) {
+                        this.replicatedModel.setParams(protoModel.params().unsafeDuplication(true));
+
+                        Nd4j.getExecutioner().commit();
+                    }
+                } else {
+                    this.replicatedModel = protoModel;
+                }
+            }
+        }
+
+        @Override
+        public void run() {
+            Nd4j.getAffinityManager().unsafeSetDevice(deviceId);
+            try {
+                // model should be replicated & initialized here
+                initializeReplicaModel();
+
+                boolean isCG = replicatedModel instanceof  ComputationGraph;
+                boolean isMLN = replicatedModel instanceof  MultiLayerNetwork;
+
+                while (shouldWork.get()) {
+                    InferenceObservable request = inputQueue.take();
+
+                    if (request != null) {
+                        counter.incrementAndGet();
+
+                        // FIXME: get rid of instanceof here, model won't change during runtime anyway
+                        if (isCG) {
+                            List<Pair<INDArray[],INDArray[]>> batches = request.getInputBatches();
+                            List<INDArray[]> out = new ArrayList<>(batches.size());
+                            try {
+                                for (Pair<INDArray[],INDArray[]> inBatch : batches) {
+                                    try {
+                                        modelLock.readLock().lock();
+                                         if(layersToOutputTo != null) {
+                                            ComputationGraph computationGraph = (ComputationGraph) replicatedModel;
+                                            INDArray[] output = computationGraph.output(Arrays.asList(layersToOutputTo),false,inBatch.getFirst(), inBatch.getSecond());
+                                            out.add(output);
+                                        } else {
+                                            INDArray[] output = ((ComputationGraph) replicatedModel).output(false, inBatch.getFirst(), inBatch.getSecond());
+                                            out.add(output);
+                                        }
+
+                                    } finally {
+                                        Nd4j.getExecutioner().commit();
+                                        modelLock.readLock().unlock();
+                                    }
+
+                                }
+                                request.setOutputBatches(out);
+                            } catch (Exception e){
+                                request.setOutputException(e);
+                            }
+                        } else if (isMLN) {
+                            List<Pair<INDArray[],INDArray[]>> batches = request.getInputBatches();
+                            List<INDArray[]> out = new ArrayList<>(batches.size());
+                            try {
+                                for (Pair<INDArray[],INDArray[]> inBatch : batches) {
+                                    INDArray f = inBatch.getFirst()[0];
+                                    INDArray fm = (inBatch.getSecond() == null ? null : inBatch.getSecond()[0]);
+                                    try {
+                                        modelLock.readLock().lock();
+                                        if(layerIndicesOutputTo != null) {
+                                            MultiLayerNetwork multiLayerNetwork = (MultiLayerNetwork) replicatedModel;
+                                            List<INDArray> indArrays = multiLayerNetwork.feedForwardToLayer(layerIndicesOutputTo[0], f, false);
+                                            out.add(new INDArray[]{indArrays.get(0)});
+                                        } else {
+                                            INDArray output = ((MultiLayerNetwork) replicatedModel).output(f, false, fm, null);
+                                            out.add(new INDArray[]{output});
+                                        }
+
+                                    } finally {
+                                        Nd4j.getExecutioner().commit();
+                                        modelLock.readLock().unlock();
+                                    }
+                                }
+                                request.setOutputBatches(out);
+                            } catch (Exception e){
+                                request.setOutputException(e);
+                            }
+                        }
+
+
+                    } else {
+                        // just do nothing, i guess and hope for next round?
+                    }
+                }
+            } catch (InterruptedException e) {
+                Thread.currentThread().interrupt();
+                // do nothing
+            } catch (Exception e) {
+                throw new RuntimeException(e);
+            } finally {
+                isStopped.set(true);
+            }
+        }
+
+        protected void shutdown() {
+            shouldWork.set(false);
+            while (!isStopped.get()) {
+                // block until main loop is finished
+            }
+        }
+    }
+
+
+    protected static class ObservablesProvider {
+        private BlockingQueue<InferenceObservable> targetQueue;
+        private long nanos;
+        private int batchLimit;
+
+        private volatile BatchedInferenceObservable currentObservable;
+        private final Object locker = new Object();
+
+        protected ObservablesProvider(long nanos, int batchLimit, @NonNull BlockingQueue<InferenceObservable> queue) {
+            this.targetQueue = queue;
+            this.nanos = nanos;
+            this.batchLimit = batchLimit;
+        }
+
+        protected InferenceObservable setInput(@NonNull Observer observer, INDArray input) {
+            return setInput(observer, new INDArray[]{input}, null);
+        }
+
+        protected InferenceObservable setInput(@NonNull Observer observer, INDArray... input) {
+            return setInput(observer, input, null);
+        }
+
+        protected InferenceObservable setInput(@NonNull Observer observer, INDArray[] input, INDArray[] inputMask) {
+            synchronized (locker) {
+                boolean isNew = false;
+                if (currentObservable == null || currentObservable.getCounter() >= batchLimit
+                        || currentObservable.isLocked()) {
+                    isNew = true;
+                    currentObservable = new BatchedInferenceObservable();
+                }
+
+                currentObservable.addInput(input, inputMask);
+                currentObservable.addObserver(observer);
+
+                try {
+                    if (isNew)
+                        targetQueue.put(currentObservable);
+                } catch (InterruptedException e) {
+                    Thread.currentThread().interrupt();
+                    throw new RuntimeException(e);
+                }
+
+                return currentObservable;
+            }
+        }
+    }
+}

deeplearning4j/deeplearning4j-parallelwrapper/src/main/java/org/deeplearning4j/parallelism/inference/InferenceMode.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.parallelism.inference;
+
+public enum InferenceMode {
+    /**
+     * input will be passed into the model as is
+     */
+    SEQUENTIAL,
+
+    /**
+     * input will be included into the batch if computation device is busy, and executed immediately otherwise
+     */
+    BATCHED,
+
+    /**
+     * Inference will applied in the calling thread instead of workers. Worker models will be using shared parameters on per-device basis.
+     */
+    INPLACE,
+}

deeplearning4j/deeplearning4j-parallelwrapper/src/main/java/org/deeplearning4j/parallelism/inference/InferenceObservable.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.parallelism.inference;
+
+import org.nd4j.linalg.api.ndarray.INDArray;
+import org.nd4j.common.primitives.Pair;
+
+import java.util.List;
+import java.util.Observer;
+
+public interface InferenceObservable {
+
+    /**
+     * Get input batches - and their associated input mask arrays, if any<br>
+     * Note that usually the returned list will be of size 1 - however, in the batched case, not all inputs
+     * can actually be batched (variable size inputs to fully convolutional net, for example). In these "can't batch"
+     * cases, multiple input batches will be returned, to be processed
+     *
+     * @return List of pairs of input arrays and input mask arrays. Input mask arrays may be null.
+     */
+    List<Pair<INDArray[],INDArray[]>> getInputBatches();
+
+    void addInput(INDArray... input);
+
+    void addInput(INDArray[] input, INDArray[] inputMasks);
+
+    void setOutputBatches(List<INDArray[]> output);
+
+    void setOutputException(Exception e);
+
+    void addObserver(Observer observer);
+
+    INDArray[] getOutput();
+}

deeplearning4j/deeplearning4j-parallelwrapper/src/main/java/org/deeplearning4j/parallelism/inference/LoadBalanceMode.java

@@ -0,0 +1,33 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.parallelism.inference;
+
+public enum LoadBalanceMode {
+    /**
+     * In this mode, `n+1 % nodes` node will be used for next request
+     */
+    ROUND_ROBIN,
+
+    /**
+     * in this mode we'll be picking free node for next request, blocking if we don't have free nodes at the moment
+     */
+    FIFO,
+}

deeplearning4j/deeplearning4j-parallelwrapper/src/main/java/org/deeplearning4j/parallelism/inference/observers/BasicInferenceObservable.java

@@ -0,0 +1,125 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.parallelism.inference.observers;
+
+import org.nd4j.shade.guava.base.Preconditions;
+import lombok.Getter;
+import lombok.NonNull;
+import lombok.extern.slf4j.Slf4j;
+import org.deeplearning4j.parallelism.inference.InferenceObservable;
+import org.nd4j.linalg.api.ndarray.INDArray;
+import org.nd4j.common.primitives.Pair;
+
+import java.util.Collections;
+import java.util.List;
+import java.util.Observable;
+
+@Slf4j
+public class BasicInferenceObservable extends Observable implements InferenceObservable {
+    private INDArray[] input;
+    private INDArray[] inputMasks;
+    @Getter
+    private long id;
+    private INDArray[] output;
+    protected Exception exception;
+    protected String[] layersToOutputTo;
+    protected int[] layerIndicesOutputTo;
+
+
+    public BasicInferenceObservable(int[] layerIndicesOutputTo,INDArray... inputs) {
+        this(layerIndicesOutputTo,inputs, null);
+    }
+
+    public BasicInferenceObservable(int[] layerIndicesOutputTo,INDArray[] inputs, INDArray[] inputMasks) {
+        super();
+        this.layerIndicesOutputTo = layerIndicesOutputTo;
+        this.input = inputs;
+        this.inputMasks = inputMasks;
+    }
+
+    public BasicInferenceObservable(String[] layersToOutputTo,INDArray... inputs) {
+        this(layersToOutputTo,inputs, null);
+    }
+
+    public BasicInferenceObservable(String[] layersToOutputTo,INDArray[] inputs, INDArray[] inputMasks) {
+        super();
+        this.layersToOutputTo = layersToOutputTo;
+        this.input = inputs;
+        this.inputMasks = inputMasks;
+    }
+
+    public BasicInferenceObservable(INDArray... inputs) {
+        this(inputs, null);
+    }
+
+    public BasicInferenceObservable(INDArray[] inputs, INDArray[] inputMasks) {
+        super();
+        this.input = inputs;
+        this.inputMasks = inputMasks;
+    }
+
+    @Override
+    public void addInput(@NonNull INDArray... input){
+        addInput(input, null);
+    }
+
+    @Override
+    public void addInput(@NonNull INDArray[] input, INDArray[] inputMasks) {
+        this.input = input;
+        this.inputMasks = inputMasks;
+    }
+
+    @Override
+    public void setOutputBatches(@NonNull List<INDArray[]> output) {
+        Preconditions.checkArgument(output.size() == 1, "Expected size 1 output: got size " + output.size());
+        this.output = output.get(0);
+        this.setChanged();
+        notifyObservers();
+    }
+
+    @Override
+    public List<Pair<INDArray[],INDArray[]>> getInputBatches() {
+        return Collections.singletonList(new Pair<>(input, inputMasks));
+    }
+
+    @Override
+    public void setOutputException(Exception exception) {
+        this.exception = exception;
+        this.setChanged();
+        notifyObservers();
+    }
+
+    @Override
+    public INDArray[] getOutput() {
+        checkOutputException();
+        return output;
+    }
+
+    protected void checkOutputException() {
+        if(exception != null) {
+            if(exception instanceof RuntimeException) {
+                throw (RuntimeException)exception;
+            } else {
+                throw new RuntimeException("Exception encountered while getting output: " + exception.getMessage(), exception);
+            }
+        }
+    }
+}

deeplearning4j/deeplearning4j-parallelwrapper/src/main/java/org/deeplearning4j/parallelism/inference/observers/BasicInferenceObserver.java

@@ -0,0 +1,51 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.parallelism.inference.observers;
+
+import lombok.extern.slf4j.Slf4j;
+
+import java.util.Observable;
+import java.util.Observer;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.locks.LockSupport;
+
+@Slf4j
+public class BasicInferenceObserver implements Observer {
+    private AtomicBoolean finished;
+
+    public BasicInferenceObserver() {
+        finished = new AtomicBoolean(false);
+    }
+
+    @Override
+    public void update(Observable o, Object arg) {
+        finished.set(true);
+    }
+
+    /**
+     * FOR DEBUGGING ONLY, TO BE REMOVED BEFORE MERGE
+     */
+    public void waitTillDone() {
+        while (!finished.get()) {
+            LockSupport.parkNanos(1000);
+        }
+    }
+}

deeplearning4j/deeplearning4j-parallelwrapper/src/main/java/org/deeplearning4j/parallelism/inference/observers/BatchedInferenceObservable.java

@@ -0,0 +1,239 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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.parallelism.inference.observers;
+
+import lombok.extern.slf4j.Slf4j;
+import org.deeplearning4j.parallelism.inference.InferenceObservable;
+import org.nd4j.linalg.api.ndarray.INDArray;
+import org.nd4j.linalg.dataset.api.DataSetUtil;
+import org.nd4j.linalg.indexing.INDArrayIndex;
+import org.nd4j.linalg.indexing.NDArrayIndex;
+import org.nd4j.common.primitives.Pair;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
+
+@Slf4j
+public class BatchedInferenceObservable extends BasicInferenceObservable implements InferenceObservable {
+    private List<INDArray[]> inputs = new ArrayList<>();
+    private List<INDArray[]> inputMasks = new ArrayList<>();
+    private List<INDArray[]> outputs = new ArrayList<>();
+    private AtomicInteger counter = new AtomicInteger(0);
+    private ThreadLocal<Integer> position = new ThreadLocal<>();
+    private List<int[]> outputBatchInputArrays = new ArrayList<>();
+
+    private final Object locker = new Object();
+
+    private ReentrantReadWriteLock realLocker = new ReentrantReadWriteLock();
+    private AtomicBoolean isLocked = new AtomicBoolean(false);
+    private AtomicBoolean isReadLocked = new AtomicBoolean(false);
+
+    public BatchedInferenceObservable() {
+
+    }
+
+    @Override
+    public void addInput(INDArray[] input, INDArray[] inputMasks) {
+        synchronized (locker) {
+            inputs.add(input);
+            this.inputMasks.add(inputMasks);
+            position.set(counter.getAndIncrement());
+
+            if (isReadLocked.get())
+                realLocker.readLock().unlock();
+        }
+    }
+
+    @Override
+    public List<Pair<INDArray[],INDArray[]>> getInputBatches() {
+        realLocker.writeLock().lock();
+        isLocked.set(true);
+
+        outputBatchInputArrays.clear();
+
+        // this method should pile individual examples into single batch
+
+        if (counter.get() > 1) {
+            int pos = 0;
+            List<Pair<INDArray[],INDArray[]>> out = new ArrayList<>();
+            int numArrays = inputs.get(0).length;
+            while(pos < inputs.size()) {
+                //First: determine which we can actually batch...
+                int lastPossible = pos;
+                for (int i = pos + 1; i < inputs.size(); i++) {
+                    if (canBatch(inputs.get(pos), inputs.get(i))) {
+                        lastPossible = i;
+                    } else {
+                        break;
+                    }
+                }
+
+                int countToMerge = lastPossible - pos + 1;
+                INDArray[][] featuresToMerge = new INDArray[countToMerge][0];
+                INDArray[][] fMasksToMerge = null;
+                int fPos = 0;
+                for( int i = pos; i <= lastPossible; i++) {
+                    featuresToMerge[fPos] = inputs.get(i);
+
+                    if(inputMasks.get(i) != null) {
+                        if(fMasksToMerge == null){
+                            fMasksToMerge = new INDArray[countToMerge][0];
+                            for( int j = 0; j < countToMerge; j++ ){
+                                fMasksToMerge[j] = null;
+                            }
+                        }
+                        fMasksToMerge[fPos] = inputMasks.get(i);
+                    }
+                    fPos++;
+                }
+
+                Pair<INDArray[],INDArray[]> merged = DataSetUtil.mergeFeatures(featuresToMerge, fMasksToMerge);
+                out.add(merged);
+
+                outputBatchInputArrays.add(new int[]{pos, lastPossible});
+                pos = lastPossible + 1;
+            }
+            realLocker.writeLock().unlock();
+            return out;
+        } else {
+            outputBatchInputArrays.add(new int[]{0,0});
+            realLocker.writeLock().unlock();
+            return Collections.singletonList(new Pair<>(inputs.get(0), inputMasks.get(0)));
+        }
+    }
+
+    private static boolean canBatch(INDArray[] first, INDArray[] candidate) {
+        //Check if we can batch these inputs into the one array. This isn't always possible - for example, some fully
+        // convolutional nets can support different input image sizes
+        //For now: let's simply require that the inputs have the same shape
+        //In the future: we'll intelligently handle the RNN variable length case
+        //Note also we can ignore input masks here - they should have shared dimensions with the input, thus if the
+        // inputs can be batched, so can the masks
+        for(int i=0; i<first.length; i++ ){
+            if(!Arrays.equals(first[i].shape(), candidate[i].shape())){
+                return false;
+            }
+        }
+        return true;
+    }
+
+    @Override
+    public void setOutputBatches(List<INDArray[]> output) {
+        //this method should split batched output INDArray[] into multiple separate INDArrays
+        int countNumInputBatches = 0;   //Counter for total number of input batches processed
+        for( int outBatchNum = 0; outBatchNum < output.size(); outBatchNum++) { //Iterate over output batch
+            INDArray[] currBatchOutputs = output.get(outBatchNum);
+            int[] inputBatchIdxs = outputBatchInputArrays.get(outBatchNum);
+            int inputBatchCount = inputBatchIdxs[1] - inputBatchIdxs[0] + 1;
+            for (int i = 0; i < inputBatchCount; i++) {
+                outputs.add(new INDArray[currBatchOutputs.length]);
+            }
+
+            // pull back results for individual input batches
+            int firstInputBatch = countNumInputBatches;
+            for (int outputNumber = 0; outputNumber < currBatchOutputs.length; outputNumber++) {    //Iterate over net outputs
+                INDArray[] split = splitExamples(currBatchOutputs[outputNumber], inputBatchIdxs[0], inputBatchIdxs[1]);
+
+                int currentInputBatch = firstInputBatch;
+                //Iterate over input batch (examples) - note that each output batch is made up of 1 or more input batches
+                for (int inputInBatch = 0; inputInBatch < inputBatchCount; inputInBatch++) {
+                    outputs.get(currentInputBatch++)[outputNumber] = split[inputInBatch];
+
+                    if(outputNumber == 0) {
+                        countNumInputBatches++;
+                    }
+                }
+            }
+        }
+
+        this.setChanged();
+        notifyObservers();
+    }
+
+    private INDArray[] splitExamples(INDArray netOutput, int firstInputComponent, int lastInputComponent){
+
+        int numSplits = lastInputComponent - firstInputComponent + 1;
+        if(numSplits == 1){
+            return new INDArray[]{netOutput};
+        } else {
+            INDArray[] out = new INDArray[numSplits];
+            INDArrayIndex[] indices = new INDArrayIndex[netOutput.rank()];
+            for(int i=1; i<indices.length; i++ ){
+                indices[i] = NDArrayIndex.all();
+            }
+            int examplesSoFar = 0;
+            for( int inNum = 0; inNum < numSplits; inNum++) {
+                var inSizeEx = inputs.get(firstInputComponent + inNum)[0].size(0);
+                indices[0] = NDArrayIndex.interval(examplesSoFar, examplesSoFar + inSizeEx);
+                out[inNum] = netOutput.get(indices);
+                examplesSoFar += inSizeEx;
+            }
+            return out;
+        }
+    }
+
+    /**
+     * PLEASE NOTE: This method is for tests only
+     *
+     * @return
+     */
+    protected List<INDArray[]> getOutputs() {
+        return outputs;
+    }
+
+    protected void setCounter(int value) {
+        counter.set(value);
+    }
+
+    public void setPosition(int pos) {
+        position.set(pos);
+    }
+
+    public int getCounter() {
+        return counter.get();
+    }
+
+
+
+    public boolean isLocked() {
+        boolean lck = !realLocker.readLock().tryLock();
+
+        boolean result = lck || isLocked.get();
+
+        if (!result)
+            isReadLocked.set(true);
+
+        return result;
+    }
+
+
+    @Override
+    public INDArray[] getOutput() {
+        // basically we should take care of splits here: each client should get its own part of output, wrt order number
+        checkOutputException();
+        return outputs.get(position.get());
+    }
+}

deeplearning4j/deeplearning4j-parallelwrapper/src/main/java9/module-info.java

@@ -0,0 +1,17 @@
+open module deeplearning4j.parallel.wrapper {
+    requires deeplearning4j.utility.iterators;
+    requires guava;
+    requires jcommander;
+    requires resources;
+    requires slf4j.api;
+    requires deeplearning4j.core;
+    requires deeplearning4j.nn;
+    requires nd4j.api;
+    requires nd4j.common;
+    exports org.deeplearning4j.parallelism;
+    exports org.deeplearning4j.parallelism.factory;
+    exports org.deeplearning4j.parallelism.inference;
+    exports org.deeplearning4j.parallelism.inference.observers;
+    exports org.deeplearning4j.parallelism.main;
+    exports org.deeplearning4j.parallelism.trainer;
+}