Package ai.djl.nn.core
Class Embedding<T>
- java.lang.Object
-
- ai.djl.nn.AbstractBlock
-
- ai.djl.nn.core.Embedding<T>
-
- Type Parameters:
T
- the type of item that should be embedded and map to the array
- All Implemented Interfaces:
Block
,AbstractEmbedding<T>
,AbstractIndexedEmbedding<T>
- Direct Known Subclasses:
TrainableWordEmbedding
public abstract class Embedding<T> extends AbstractBlock implements AbstractIndexedEmbedding<T>
An Embedding block map a collection of items to 1-Dimensional representativeNDArray
s.
-
-
Nested Class Summary
Nested Classes Modifier and Type Class Description static class
Embedding.BaseBuilder<T,B extends Embedding.BaseBuilder<T,B>>
protected class
Embedding.DefaultEmbedding
protected class
Embedding.DefaultItem
-
Field Summary
Fields Modifier and Type Field Description protected Parameter
embedding
protected int
embeddingSize
protected AbstractIndexedEmbedding<T>
fallthroughEmbedding
protected int
numEmbeddings
protected SparseFormat
sparseFormat
-
Fields inherited from class ai.djl.nn.AbstractBlock
children, inputNames, inputShapes, parameters, version
-
-
Constructor Summary
Constructors Modifier Constructor Description Embedding(NDArray embedding)
Constructs a pretrained embedding.Embedding(NDArray embedding, SparseFormat format)
Constructs a pretrained embedding.protected
Embedding(Embedding.BaseBuilder<T,?> baseBuilder)
-
Method Summary
All Methods Static Methods Instance Methods Concrete Methods Modifier and Type Method Description NDArray
embed(NDManager manager, T[] items)
Embeds an array of items.static NDList
embedding(NDArray input, NDArray weight, SparseFormat sparse)
A simple lookup table that looks up embeddings in a fixed dictionary and size.protected NDList
forwardInternal(ParameterStore parameterStore, 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.void
loadParameters(NDManager manager, java.io.DataInputStream is)
Loads the parameters from the given input stream.void
prepare(Shape[] inputShapes)
Sets the shape ofParameter
s.void
saveParameters(java.io.DataOutputStream os)
Writes the parameters of the block to the given outputStream.-
Methods inherited from class ai.djl.nn.AbstractBlock
addChildBlock, addParameter, beforeInitialize, cast, clear, describeInput, forward, forward, forwardInternal, getChildren, getDirectParameters, getParameters, initialize, initializeChildBlocks, isInitialized, loadMetadata, readInputShapes, saveInputShapes, saveMetadata, 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.core.AbstractEmbedding
hasItem
-
Methods inherited from interface ai.djl.nn.core.AbstractIndexedEmbedding
decode, embed, encode, unembed
-
-
-
-
Field Detail
-
numEmbeddings
protected int numEmbeddings
-
embeddingSize
protected int embeddingSize
-
sparseFormat
protected SparseFormat sparseFormat
-
fallthroughEmbedding
protected AbstractIndexedEmbedding<T> fallthroughEmbedding
-
embedding
protected Parameter embedding
-
-
Constructor Detail
-
Embedding
protected Embedding(Embedding.BaseBuilder<T,?> baseBuilder)
-
Embedding
public Embedding(NDArray embedding)
Constructs a pretrained embedding.- Parameters:
embedding
- the embedding array
-
Embedding
public Embedding(NDArray embedding, SparseFormat format)
Constructs a pretrained embedding.- Parameters:
embedding
- the embedding arrayformat
- whether to compute row sparse gradient in the backward calculation
-
-
Method Detail
-
prepare
public void prepare(Shape[] inputShapes)
Sets the shape ofParameter
s.- Overrides:
prepare
in classAbstractBlock
- Parameters:
inputShapes
- the shapes of inputs
-
getOutputShapes
public Shape[] getOutputShapes(Shape[] inputShapes)
Returns the expected output shapes of the block for the specified input shapes.- Specified by:
getOutputShapes
in interfaceBlock
- Parameters:
inputShapes
- the shapes of the inputs- Returns:
- the expected output shapes of the block
-
forwardInternal
protected NDList forwardInternal(ParameterStore parameterStore, 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:
parameterStore
- the parameter storeinputs
- the input NDListtraining
- true for a training forward passparams
- optional parameters- Returns:
- the output of the forward pass
-
saveParameters
public void saveParameters(java.io.DataOutputStream os) throws java.io.IOException
Writes the parameters of the block to the given outputStream.- Specified by:
saveParameters
in interfaceBlock
- Overrides:
saveParameters
in classAbstractBlock
- Parameters:
os
- the outputstream to save the parameters to- Throws:
java.io.IOException
- if an I/O error occurs
-
loadParameters
public void loadParameters(NDManager manager, java.io.DataInputStream is) throws java.io.IOException, MalformedModelException
Loads the parameters from the given input stream.- Specified by:
loadParameters
in interfaceBlock
- Overrides:
loadParameters
in classAbstractBlock
- Parameters:
manager
- an NDManager to create the parameter arraysis
- the inputstream that stream the parameter values- Throws:
java.io.IOException
- if an I/O error occursMalformedModelException
- if the model file is corrupted or unsupported
-
embed
public NDArray embed(NDManager manager, T[] items)
Embeds an array of items.- Specified by:
embed
in interfaceAbstractEmbedding<T>
- Parameters:
manager
- the manager for the new embeddingsitems
- the items to embed- Returns:
- the embedding
NDArray
of Shape(items.length, embeddingSize)
-
embedding
public static NDList embedding(NDArray input, NDArray weight, SparseFormat sparse)
A simple lookup table that looks up embeddings in a fixed dictionary and size.- Parameters:
input
- NDArray containing indices into the embedding matrixweight
- The embedding matrix with number of rows equal to the maximum possible index + 1, and number of columns equal to the embedding sizesparse
- SparseFormat of the gradient- Returns:
- output NDArray
-
-