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