Spark API Reference
API reference for SparkDl4jMultiLayer, SparkComputationGraph, and TrainingMaster
This page documents the key classes for distributed training with DL4J on Spark. For setup and how-to guides, see the Spark How-To guide. For an introduction to the architecture, see the Distributed Training Overview.
Contents:
SparkDl4jMultiLayer
Main class for training MultiLayerNetwork networks using Spark. Also supports distributed evaluation and inference.
Constructor
SparkDl4jMultiLayer(JavaSparkContext sc, MultiLayerConfiguration conf, TrainingMaster trainingMaster)
SparkDl4jMultiLayer(JavaSparkContext sc, MultiLayerNetwork network, TrainingMaster trainingMaster)trainingMaster may be null when the instance is used only for evaluation or inference (not training).
Network Access
public MultiLayerNetwork getNetwork()
public void setNetwork(MultiLayerNetwork network)
public JavaSparkContext getSparkContext()
public TrainingMaster getTrainingMaster()Training
Train from an RDD of DataSet objects. Note: fitting directly from RDD<DataSet> is not the recommended approach — prefer saving data to disk and using fit(String).
Train from a directory of serialized DataSet objects on network storage (HDFS, S3, etc.). The directory must contain files serialized using DataSet.save(OutputStream). This is the preferred training method.
Train from an RDD of paths pointing to serialized DataSet objects.
Train from an RDD of paths using a custom DataSetLoader to deserialize each file.
Convenience methods for compatibility with Spark MLLib LabeledPoint format. fitContinuousLabeledPoint is for regression targets.
Scoring
Returns the average minibatch loss from the most recent fit call, averaged across all workers.
Calculate the total or average loss across an entire RDD. minibatchSize controls memory use during scoring; default is DEFAULT_EVAL_SCORE_BATCH_SIZE.
Return a per-example loss. Unlike calculateScore, this returns one value per example (not an aggregate).
Evaluation
Classification metrics: accuracy, F1, precision, recall. evalNumWorkers controls how many network copies are used per Spark executor (reduces memory usage for large networks). Default is DEFAULT_EVAL_WORKERS.
ROC curve evaluation for single-output binary classifiers.
ROC evaluation for multi-class classifiers (one ROC curve per class).
Regression metrics: MSE, MAE, R2, etc.
Perform multiple evaluations in a single pass over the data — more efficient than calling evaluation methods sequentially.
Example:
Distributed Inference
Run inference on a keyed RDD of feature arrays. Returns a keyed RDD of predictions. The key K is used to associate inputs with outputs (Spark RDDs are unordered). Does not support mask arrays.
Overload that accepts an input mask array (for variable-length sequences).
Statistics and Debugging
Enable/disable detailed training statistics collection. Disabled by default. When enabled, requires internet access to an NTP server unless the time source is overridden (see troubleshooting guide).
Get/set the default number of network instances used for distributed evaluation per executor. Setting this lower than the number of Spark threads per executor reduces memory consumption for large models.
SparkComputationGraph
Main class for training ComputationGraph networks using Spark. Mirrors SparkDl4jMultiLayer but supports multi-input/multi-output networks via MultiDataSet.
Constructor
Network Access
Training
Training methods mirror SparkDl4jMultiLayer. The fitMultiDataSet and fitPathsMultiDataSet variants accept MultiDataSet objects, enabling multi-input/multi-output training.
Scoring
Evaluation
Distributed Inference
Returns INDArray[] per example (one array per output node) rather than a single INDArray.
Evaluation Workers
SharedTrainingMaster
Implements distributed training using the Strom 2015 compressed gradient sharing algorithm. This is the recommended TrainingMaster implementation.
Serialization
Serialize/deserialize the configuration. Useful for saving the training configuration alongside saved models.
Builder
Core Training Parameters
Minibatch size on each worker. The source RDD DataSets may have a different size — DL4J will split or combine them as needed.
Number of training threads per cluster node. Default: -1 (auto-detect based on hardware). On GPU nodes, set to the number of GPUs. On CPU nodes, typically 1; for machines with many cores and large core counts, you may increase this (set OMP_NUM_THREADS accordingly to avoid over-subscription).
Threshold and Residual Configuration
Algorithm that determines the gradient encoding threshold. Default: AdaptiveThresholdAlgorithm which adjusts the threshold to keep sparsity between 0.0001 and 0.01. See Spark How-To: Encoding Thresholds for details.
Deprecated. Use thresholdAlgorithm(new FixedThresholdAlgorithm(value)) instead.
Controls how the residual vector (un-communicated gradient accumulation) is post-processed. Default: ResidualClippingPostProcessor(5.0, 5) — clips the residual to 5x the threshold every 5 steps, preventing residual explosion.
Cluster Topology
Communication topology. Options:
MeshBuildMode.PLAIN: Master relays all updates. Suitable for clusters with fewer than ~32 nodes.MeshBuildMode.MESH: Non-binary tree topology. Reduces master load. Recommended for larger clusters.
Data Handling
How to handle RDD<DataSet> training data:
RDDTrainingApproach.Export(default): exports to temporary HDFS directory before training.RDDTrainingApproach.Direct: uses data directly from the RDD.
Prefer Export — it avoids redundant recomputation and is more memory-efficient.
Base directory for temporary data export when using RDDTrainingApproach.Export. Default: {hadoop.tmp.dir}/dl4j/.
Storage level for RDD<DataSet> persistence when using RDDTrainingApproach.Direct. Default: MEMORY_ONLY_SER. See caching guidance — never use MEMORY_ONLY with DL4J RDDs.
Controls how data is repartitioned before training. Default: DefaultRepartitioner (equalizes up to 5000 partitions). Imbalanced partitions cause "end-of-epoch" stalls where the cluster waits for the slowest partition.
Worker Configuration
Number of minibatches to asynchronously prefetch on each worker. Default: 2. Increase if ETL is a bottleneck; reduce if memory is tight.
Configure periodic garbage collection on workers. Default (1.0.0-beta3+): GC every 5000 ms. Disable or increase the interval when using workspaces to avoid unnecessary GC pauses.
Debugging
When enabled, logs threshold, sparsity ratio, and encoding statistics on each worker at each iteration. Useful for diagnosing threshold issues. Has performance overhead — use only during investigation.
Enable Spark-level training statistics collection. Disabled by default.
Artificially extends each iteration by sleeping for timeMs milliseconds. For debugging only — never use in production.
Miscellaneous
RNG seed for repeatable data partitioning.
Custom Aeron transport implementation. Not required for standard UDP communication.
ParameterAveragingTrainingMaster
Synchronous SGD implementation via Spark. Workers train independently for averagingFrequency minibatches, then parameters are averaged on the master. Superseded by SharedTrainingMaster — prefer gradient sharing for new projects.
Serialization
Builder
rddDataSetNumExamples is the number of examples per DataSet object in the source RDD.
Core Parameters
Minibatch size per worker per averaging step.
How often (in number of minibatches) workers synchronize with the master. Too low (e.g., 1) creates excessive network traffic. Too high (e.g., > 20) can hurt convergence. A value of 5–10 is a reasonable starting point.
Depth of the aggregation tree used to reduce parameters back to the master. Default: 2. Increase for large clusters with many partitions to avoid the driver becoming a bottleneck.
Whether to include the optimizer state (momentum buffers, AdaGrad accumulators, etc.) in the averaged parameters. Default: true. Setting to false doubles or more the effective parameter server bandwidth but disables updater state sharing, which may harm convergence for adaptive optimizers.
Data Handling
Number of minibatches to asynchronously prefetch on each worker. Default: 0 (no prefetching).
When to repartition training data (default: always repartition to ensure balanced partitions). Values: Always, Never, NumPartitionsWorkersDiffers.
How to repartition. SparkDefault uses Spark's built-in shuffle; Balanced balances the number of examples per partition (not just the number of partitions).
Storage level for RDD<DataSet> persistence. Default: MEMORY_ONLY_SER. See caching guidance.
Storage level for path-based data (PortableDataStream RDDs from fit(String) or fitPaths). Default: MEMORY_ONLY.
Same semantics as in SharedTrainingMaster.Builder.
Miscellaneous
Training Hook Interface
TrainingHook instances receive callbacks before and after each training step on workers. Can be used for custom monitoring or parameter manipulation.
VoidConfiguration
VoidConfiguration is a required companion to SharedTrainingMaster that configures the Aeron-based communication layer.
unicastPort: Any available UDP port. Must be open (both inbound and outbound) on all cluster nodes. Configure your firewall/security groups accordingly.
networkMask: CIDR-format network mask that selects the network interface used for Aeron communication. Required when running on YARN or in environments (AWS, Azure) where Spark's detected IP may differ from the desired communication interface. Example: 192.168.0.0/16, 10.1.2.0/24.
controllerAddress: The IP address of the Spark master/driver. Workers use this to connect to the parameter server master.
As a fallback when automatic interface selection fails, set the DL4J_VOID_IP environment variable on each node to the IP address to use for Aeron communication.
Last updated
Was this helpful?