ai.djl.nn.transformer

Class TransformerEncoderBlock

java.lang.Object

ai.djl.nn.AbstractBlock

ai.djl.nn.transformer.TransformerEncoderBlock

All Implemented Interfaces:

Block

public class TransformerEncoderBlock
extends AbstractBlock

Self-Attention based transformer encoder block.

Field Summary

Fields inherited from class ai.djl.nn.AbstractBlock

children, inputNames, inputShapes, parameters, version

Constructor Summary

Constructors

Constructor and Description
TransformerEncoderBlock(int embeddingSize, int headCount, int hiddenSize, float dropoutProbability, java.util.function.Function<NDList,NDList> activationFunction) Creates a transformer encoder block.

Method Summary

Modifier and Type	Method and Description
protected NDList	forwardInternal(ParameterStore ps, NDList inputs, boolean training, ai.djl.util.PairList<java.lang.String,java.lang.Object> params) A helper for Block.forward(ParameterStore, NDList, boolean, PairList) after initialization.
Shape[]	getOutputShapes(Shape[] inputShapes) Returns the expected output shapes of the block for the specified input shapes.
void	initializeChildBlocks(NDManager manager, DataType dataType, Shape... inputShapes) Initializes the Child blocks of this block.

Methods inherited from class ai.djl.nn.AbstractBlock

addChildBlock, addParameter, beforeInitialize, cast, clear, describeInput, forward, forward, forwardInternal, getChildren, getDirectParameters, getParameters, initialize, isInitialized, loadMetadata, loadParameters, prepare, readInputShapes, saveInputShapes, saveMetadata, saveParameters, setInitializer, setInitializer, setInitializer, toString

Methods inherited from class java.lang.Object

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

Methods inherited from interface ai.djl.nn.Block

forward, validateLayout

Constructor Detail

TransformerEncoderBlock

public TransformerEncoderBlock(int embeddingSize,
                               int headCount,
                               int hiddenSize,
                               float dropoutProbability,
                               java.util.function.Function<NDList,NDList> activationFunction)

Creates a transformer encoder block.

Parameters:

embeddingSize - the embedding size for tokens

headCount - number of attention blocks

hiddenSize - the hidden size for fully connected networks

dropoutProbability - dropout probability

activationFunction - activation function

Method Detail

getOutputShapes

public Shape[] getOutputShapes(Shape[] inputShapes)

Returns the expected output shapes of the block for the specified input shapes.

Parameters:

inputShapes - the shapes of the inputs

Returns:

the expected output shapes of the block

initializeChildBlocks

public void initializeChildBlocks(NDManager manager,
                                  DataType dataType,
                                  Shape... inputShapes)

Initializes the Child blocks of this block. You need to override this method if your subclass has child blocks. Used to determine the correct input shapes for child blocks based on the requested input shape for this block.

Overrides:

initializeChildBlocks in class AbstractBlock

Parameters:

manager - the manager to use for initialization

dataType - the requested data type

inputShapes - the expected input shapes for this block

forwardInternal

protected NDList forwardInternal(ParameterStore ps,
                                 NDList inputs,
                                 boolean training,
                                 ai.djl.util.PairList<java.lang.String,java.lang.Object> params)

A helper for Block.forward(ParameterStore, NDList, boolean, PairList) after initialization.

Specified by:

forwardInternal in class AbstractBlock

Parameters:

ps - the parameter store

inputs - the input NDList

training - true for a training forward pass

params - optional parameters

Returns:

the output of the forward pass