Package ai.djl.nn.transformer
Class IdEmbedding
- java.lang.Object
-
- ai.djl.nn.AbstractBlock
-
- ai.djl.nn.transformer.IdEmbedding
-
- All Implemented Interfaces:
Block
public final class IdEmbedding extends AbstractBlock
An Embedding from integer ids to float vectors. Output shape is the input shape + one dimension for the embedding. E.g. If input shape is (-1, 128), embedding size is 1024, then the output shape is (-1, 128, 1024)
-
-
Nested Class Summary
Nested Classes Modifier and Type Class Description static class
IdEmbedding.Builder
The Builder to construct anIdEmbedding
type ofBlock
.
-
Field Summary
-
Fields inherited from class ai.djl.nn.AbstractBlock
children, inputNames, inputShapes, parameters, version
-
-
Method Summary
All Methods Instance Methods Concrete Methods Modifier and Type Method Description protected NDList
forwardInternal(ParameterStore ps, NDList inputs, boolean training, ai.djl.util.PairList<java.lang.String,java.lang.Object> params)
A helper forBlock.forward(ParameterStore, NDList, boolean, PairList)
after initialization.Shape[]
getOutputShapes(Shape[] inputShapes)
Returns the expected output shapes of the block for the specified input shapes.NDArray
getValue(ParameterStore ps, Device device, boolean training)
Quick hack for bert model to acces embedding table, replace by a proper function to calculate raw logits from embeddings.void
initializeChildBlocks(NDManager manager, DataType dataType, Shape... inputShapes)
Initializes the Child blocks of this block.NDArray
probabilities(ParameterStore parameterStore, NDArray input, boolean training)
Turns an array of embeddings of shape (d0 ...-
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
-
-
-
-
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
-
forwardInternal
protected NDList forwardInternal(ParameterStore ps, NDList inputs, boolean training, ai.djl.util.PairList<java.lang.String,java.lang.Object> params)
A helper forBlock.forward(ParameterStore, NDList, boolean, PairList)
after initialization.- Specified by:
forwardInternal
in classAbstractBlock
- Parameters:
ps
- the parameter storeinputs
- the input NDListtraining
- true for a training forward passparams
- optional parameters- Returns:
- the output of the forward pass
-
probabilities
public NDArray probabilities(ParameterStore parameterStore, NDArray input, boolean training)
Turns an array of embeddings of shape (d0 ... dN, E) into an array of log probabilities of shape (d0 ... dN, D) that shows the probability distribution that a given embedding corresponds to an entry in the internal embedding table.- Parameters:
parameterStore
- the parameters storeinput
- the embeddings to create log probabilities fortraining
- true for a training forward pass- Returns:
- log probabilities for each embedding
-
getValue
public NDArray getValue(ParameterStore ps, Device device, boolean training)
Quick hack for bert model to acces embedding table, replace by a proper function to calculate raw logits from embeddings. TODO: replace by function to get logits- Parameters:
ps
- the parameter storedevice
- device to get internal table fortraining
- true for a training forward pass- Returns:
- this embedding table as an array on the given device
-
initializeChildBlocks
public void initializeChildBlocks(NDManager manager, DataType dataType, Shape... inputShapes)
Description copied from class:AbstractBlock
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 classAbstractBlock
- Parameters:
manager
- the manager to use for initializationdataType
- the requested data typeinputShapes
- the expected input shapes for this block
-
-