ai.djl.training

Class ParallelTrain

java.lang.Object

ai.djl.training.ParallelTrain

public final class ParallelTrain
extends java.lang.Object

Helper for easy training of a whole model on multiple GPUs in parallel.

Constructor Summary

Constructors

Constructor and Description
ParallelTrain(Device[] devices) Build a ParallelTrain for the given devices.

Method Summary

Modifier and Type	Method and Description
void	trainBatch(Trainer trainer, Batch batch) Trains the model with one iteration of the given Batch of data.

Methods inherited from class java.lang.Object

clone, equals, finalize, getClass, hashCode, notify, notifyAll, toString, wait, wait, wait

Constructor Detail

ParallelTrain

public ParallelTrain(Device[] devices)

Build a ParallelTrain for the given devices.

Parameters:

devices - the devices to train on.

Method Detail

trainBatch

public void trainBatch(Trainer trainer,
                       Batch batch)

Trains the model with one iteration of the given Batch of data.

Parameters:

trainer - the trainer to validate the batch with

batch - a Batch that contains data, and its respective labels

Throws:

java.lang.IllegalArgumentException - if the batch engine does not match the trainer engine