Package ai.djl.translate

Class PaddingStackBatchifier

java.lang.Object
ai.djl.translate.PaddingStackBatchifier
All Implemented Interfaces:
Batchifier, Serializable
public final class PaddingStackBatchifier extends Object implements Batchifier
The padding stack batchifier is a StackBatchifier that also pads elements to reach the same length.
See Also:

  • Method Details

    • batchify

      public NDList batchify(NDList[] inputs)
      Converts an array of Record NDLists into a combined Batch NDList.

      The size of the input array is the batch size. The data in each of the NDList are assumed to be the same, and are batched together to form one batched NDList.

      Specified by:
      batchify in interface Batchifier
      Parameters:
      inputs - the input array of NDList where each element is a
      Returns:
      the batchified NDList

    • unbatchify

      public NDList[] unbatchify(NDList inputs)
      Splits a combined Batch NDList into it's constituent Record NDLists.

      This reverses the batchify operation.

      Specified by:
      unbatchify in interface Batchifier
      Parameters:
      inputs - the NDList that needs to be 'unbatchified'
      Returns:
      an array of NDLists, of size equal to batch size, where each NDList is one element from the batch of inputs

    • split

      public NDList[] split(NDList list, int numOfSlices, boolean evenSplit)
      Partitions the given Batch NDList into multiple Batch lists with smaller batch size.

      As an example, this function might be used for multi-GPU training where it takes the main batch and splits it into sub-batches that can be run on each GPU.

      This function unbatchifies the input NDList, redistributes them into the given number of slices, and then batchify each of the slices to form an array of NDList.

      Specified by:
      split in interface Batchifier
      Parameters:
      list - the NDList that needs to be split
      numOfSlices - the number of slices the list must be sliced into
      evenSplit - whether each slice must have the same shape
      Returns:
      an array of NDList that contains all the slices

    • findMaxSize

      public static long findMaxSize(NDList[] inputs, int arrayIndex, int dimIndex)
      Finds the maximum size for a particular array/dimension in a batch of inputs (which can be padded to equalize their sizes).
      Parameters:
      inputs - the batch of inputs
      arrayIndex - the array (for each NDList in the batch)
      dimIndex - for the array in each NDList in the batch
      Returns:
      the maximum size

    • padArrays

      public static long[] padArrays(NDList[] inputs, int arrayIndex, int dimIndex, NDArray padding, long maxSize)
      Pads the arrays at a particular dimension to all have the same size (updating inputs in place).
      Parameters:
      inputs - the batch of inputs
      arrayIndex - the array (for each NDList in the batch)
      dimIndex - for the array in each NDList in the batch
      padding - the padding to use. Say you have a batch of arrays of Shape(10, ?, 3) and you are padding the "?" dimension. There are two padding modes:
      • If you give padding of Shape(1, 3) (same dimensionality as required), it will be repeated with NDArray.repeat(long) as necessary
      • If you give padding of Shape(3) or Shape(0) (smaller dimensionality as required), it will be broadcasted with NDArray.broadcast(Shape) to reach the full required Shape(?, 3)
      maxSize - the size that each array will be padded to in that dimension. In the example above, the padding to be applied to the "?" dimension.
      Returns:
      the original valid length for each dimension in the batch (same length as inputs.length). The inputs will be updated in place.

    • builder

      public static PaddingStackBatchifier.Builder builder()
      Returns a PaddingStackBatchifier.Builder.
      Returns:
      a PaddingStackBatchifier.Builder