LoadAllTPUEmbeddingParameters

public final class LoadAllTPUEmbeddingParameters

An op that loads optimization parameters into embedding memory.

An op that loads optimization parameters into embedding memory. Must be preceded by a ConfigureTPUEmbeddingHost op that sets up the correct embedding table configuration. For example, this op is used to install parameters that are loaded from a checkpoint before a training loop is executed. For Adagrad, auxiliary1 should be the accumulators. For SGD, all of the auxiliary* values should be empty. For FTRL, auxiliary1 should be the accumulators and auxiliary2 should be the linear terms. For ADAM, auxiliary1 should be the momenta and auxiliary2 should be the velocities.

Public Methods

static LoadAllTPUEmbeddingParameters
create(Scope scope, Iterable<Operand<Float>> parameters, Iterable<Operand<Float>> auxiliary1, Iterable<Operand<Float>> auxiliary2, Iterable<Operand<Float>> auxiliary3, Iterable<Operand<Float>> auxiliary4, Iterable<Operand<Float>> auxiliary5, Iterable<Operand<Float>> auxiliary6, Iterable<Operand<Float>> auxiliary7, String config, Long numShards, Long shardId)
Factory method to create a class wrapping a new LoadAllTPUEmbeddingParameters operation.

Inherited Methods

Public Methods

public static LoadAllTPUEmbeddingParameters create (Scope scope, Iterable<Operand<Float>> parameters, Iterable<Operand<Float>> auxiliary1, Iterable<Operand<Float>> auxiliary2, Iterable<Operand<Float>> auxiliary3, Iterable<Operand<Float>> auxiliary4, Iterable<Operand<Float>> auxiliary5, Iterable<Operand<Float>> auxiliary6, Iterable<Operand<Float>> auxiliary7, String config, Long numShards, Long shardId)

Factory method to create a class wrapping a new LoadAllTPUEmbeddingParameters operation.

Parameters
scope current scope
parameters A list of tensors, one for each embedding table, containing the initial embedding table parameters to use in embedding lookups.
auxiliary1 A list of tensors, one for each embedding table, containing the initial values of the first auxiliary optimization parameter to use in embedding training loop updates. The shape of each entry is ignored (and thus can be empty) for those tables whose optimization algorithms do not have at least one auxiliary parameter.
auxiliary2 A list of tensors, one for each embedding table, containing the initial values of the second auxiliary optimization parameter to use in embedding training loop updates. The shape of each entry is ignored (and thus can be empty) for those tables whose optimization algorithms do not have at least two auxiliary
auxiliary3 A list of tensors, one for each embedding table, containing the initial values of the third auxiliary optimization parameter to use in embedding training loop updates. The shape of each entry is ignored (and thus can be empty) for those tables whose optimization algorithms do not have three auxiliary parameters.
auxiliary4 A list of tensors, one for each embedding table, containing the initial values of the second auxiliary optimization parameter to use in embedding training loop updates. The shape of each entry is ignored (and thus can be empty) for those tables whose optimization algorithms do not have at least four auxiliary
auxiliary5 A list of tensors, one for each embedding table, containing the initial values of the third auxiliary optimization parameter to use in embedding training loop updates. The shape of each entry is ignored (and thus can be empty) for those tables whose optimization algorithms do not have five auxiliary parameters.
auxiliary6 A list of tensors, one for each embedding table, containing the initial values of the second auxiliary optimization parameter to use in embedding training loop updates. The shape of each entry is ignored (and thus can be empty) for those tables whose optimization algorithms do not have at least six auxiliary
auxiliary7 A list of tensors, one for each embedding table, containing the initial values of the third auxiliary optimization parameter to use in embedding training loop updates. The shape of each entry is ignored (and thus can be empty) for those tables whose optimization algorithms do not have sevan auxiliary parameters.
config An TPUEmbeddingConfiguration proto describing the table parameters being loaded, serialized to a string.
numShards Number of shards into which the embedding tables are divided.
shardId Identifier of shard for this operation.
Returns
  • a new instance of LoadAllTPUEmbeddingParameters