tfa.seq2seq.GreedyEmbeddingSampler

A inference sampler that takes the maximum from the output distribution.

Inherits From: Sampler

tfa.seq2seq.GreedyEmbeddingSampler(
    embedding_fn: Optional[Callable] = None
)

Used in the notebooks

Used in the tutorials
TensorFlow Addons Networks : Sequence-to-Sequence NMT with Attention Mechanism

Uses the argmax of the output (treated as logits) and passes the result through an embedding layer to get the next input.

Args
embedding_fn	A optional callable that takes a vector tensor of ids (argmax ids). The returned tensor will be passed to the decoder input. Default to use tf.nn.embedding_lookup.

Attributes
batch_size	Batch size of tensor returned by sample. Returns a scalar int32 tensor. The return value might not available before the invocation of initialize(), in this case, ValueError is raised.
sample_ids_dtype	DType of tensor returned by sample. Returns a DType. The return value might not available before the invocation of initialize().
sample_ids_shape	Shape of tensor returned by sample, excluding the batch dimension. Returns a TensorShape. The return value might not available before the invocation of initialize().

Methods

initialize

View source

initialize(
    embedding, start_tokens=None, end_token=None
)

Initialize the GreedyEmbeddingSampler.

Args
embedding	tensor that contains embedding states matrix. It will be used to generate generate outputs with start_tokens and end_token. The embedding will be ignored if the embedding_fn has been provided at init().
start_tokens	int32 vector shaped [batch_size], the start tokens.
end_token	int32 scalar, the token that marks end of decoding.

Returns
Tuple of two items: (finished, self.start_inputs).

Raises
ValueError	if start_tokens is not a 1D tensor or end_token is not a scalar.

next_inputs

View source

next_inputs(
    time, outputs, state, sample_ids
)

next_inputs_fn for GreedyEmbeddingHelper.

sample

View source

sample(
    time, outputs, state
)

sample for GreedyEmbeddingHelper.