ai.djl.training.dataset

Class Batch

java.lang.Object

ai.djl.training.dataset.Batch

All Implemented Interfaces:

java.lang.AutoCloseable

public class Batch
extends java.lang.Object
implements java.lang.AutoCloseable

A Batch is used to hold multiple items (data and label pairs) from a Dataset.

When training and performing inference, it is often more efficient to run multiple items through a network simultaneously rather than one at a time. For this reason, much of the API is oriented around the Batch class.

In a Batch, data and label are each an NDList. The data NDList represents the data for each input in the batch. The number of NDArrays in the NDList is based on the number of different kinds of inputs, not the batch size. Similarly, the label NDList represents the labels for each kind of output.

For example, an Image Question and Answer dataset has two inputs: an image and a question. In this case, the data in the Batch will be an NDList containing an NCHW image NDArray and an NTC question NDArray. The label will be an NDList containing only an NTC answer NDArray.

In order to differentiate a batch vs a single record (despite them both consisting of two NDLists), we have the Batch and the Record respectively.

Constructor Summary

Constructors

Constructor and Description
Batch(NDManager manager, NDList data, NDList labels) Creates a new instance of Batch with the given manager, data and labels.
Batch(NDManager manager, NDList data, NDList labels, Batchifier batchifier) Creates a new instance of Batch with the given manager, data and labels.

Method Summary

Modifier and Type	Method and Description
void	close()
NDList	getData() Gets the data of this Batch.
NDList	getLabels() Gets the labels corresponding to the data of this Batch.
NDManager	getManager() Gets the NDManager that is attached to this Batch.
Batch[]	split(Device[] devices, boolean evenSplit) Splits the data and labels in the Batch across the given devices.

Methods inherited from class java.lang.Object

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

Constructor Detail

Batch

public Batch(NDManager manager,
             NDList data,
             NDList labels)

Creates a new instance of Batch with the given manager, data and labels.

Parameters:

manager - the NDManager to be attached to data and labels

data - the NDList containing the data

labels - the NDList containing the labels

Batch

public Batch(NDManager manager,
             NDList data,
             NDList labels,
             Batchifier batchifier)

Creates a new instance of Batch with the given manager, data and labels.

Parameters:

manager - the NDManager to be attached to data and labels

data - the NDList containing the data

labels - the NDList containing the labels

batchifier - the Batchifier that is used for split

Method Detail

getManager

public NDManager getManager()

Gets the NDManager that is attached to this Batch.

Returns:

the NDManager attached to this Batch

getData

public NDList getData()

Gets the data of this Batch.

Returns:

an NDList that contains the data

getLabels

public NDList getLabels()

Gets the labels corresponding to the data of this Batch.

Returns:

an NDList that contains the labels

close

public void close()

Specified by:

close in interface java.lang.AutoCloseable

split

public Batch[] split(Device[] devices,
                     boolean evenSplit)

Splits the data and labels in the Batch across the given devices.

if evenSplit is false, that last device may have a smaller batch than the rest.

Parameters:

devices - an array of Device across which the data must be split

evenSplit - whether each slice must have the same shape

Returns:

an array of Batch, each of which corresponds to a Device