Package ai.djl.nn.core

Class Embedding<T>

java.lang.Object
ai.djl.nn.AbstractBaseBlock
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 representative NDArrays.

  • Field Details

    • numEmbeddings

      protected int numEmbeddings

    • embeddingSize

      protected int embeddingSize

    • sparseFormat

      protected SparseFormat sparseFormat

    • fallthroughEmbedding

      protected AbstractIndexedEmbedding<T> fallthroughEmbedding

    • embedding

      protected Parameter embedding

  • Constructor Details

    • Embedding

      protected Embedding(Embedding.BaseBuilder<T,?> baseBuilder)

    • Embedding

      protected Embedding(NDArray embedding)
      Constructs a pretrained embedding.
      Parameters:
      embedding - the embedding array

    • Embedding

      protected Embedding(NDArray embedding, SparseFormat format)
      Constructs a pretrained embedding.

      Because it is created with preTrained data, it is created as a frozen block. If you with to update it, call Block.freezeParameters(boolean).

      Parameters:
      embedding - the embedding array
      format - whether to compute row sparse gradient in the backward calculation

  • Method Details

    • prepare

      public void prepare(Shape[] inputShapes)
      Sets the shape of Parameters.
      Overrides:
      prepare in class AbstractBaseBlock
      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 interface Block
      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<String,Object> params)
      A helper for Block.forward(ParameterStore, NDList, boolean, PairList) after initialization.
      Specified by:
      forwardInternal in class AbstractBaseBlock
      Parameters:
      parameterStore - the parameter store
      inputs - the input NDList
      training - true for a training forward pass
      params - optional parameters
      Returns:
      the output of the forward pass

    • saveParameters

      public void saveParameters(DataOutputStream os) throws IOException
      Writes the parameters of the block to the given outputStream.
      Specified by:
      saveParameters in interface Block
      Overrides:
      saveParameters in class AbstractBaseBlock
      Parameters:
      os - the outputstream to save the parameters to
      Throws:
      IOException - if an I/O error occurs

    • loadParameters

      public void loadParameters(NDManager manager, DataInputStream is) throws IOException, MalformedModelException
      Loads the parameters from the given input stream.
      Specified by:
      loadParameters in interface Block
      Overrides:
      loadParameters in class AbstractBaseBlock
      Parameters:
      manager - an NDManager to create the parameter arrays
      is - the inputstream that stream the parameter values
      Throws:
      IOException - if an I/O error occurs
      MalformedModelException - 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 interface AbstractEmbedding<T>
      Parameters:
      manager - the manager for the new embeddings
      items - 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 matrix
      weight - The embedding matrix with number of rows equal to the maximum possible index + 1, and number of columns equal to the embedding size
      sparse - SparseFormat of the gradient
      Returns:
      output NDArray