def forward(
self, # type: ignore
ids: Optional[torch.LongTensor] = None,
soft_ids: Optional[torch.Tensor] = None,
**kwargs) -> torch.Tensor:
r"""Embeds (soft) ids.
Either :attr:`ids` or :attr:`soft_ids` must be given, and they
must not be given at the same time.
Args:
ids (optional): An integer tensor containing the ids to embed.
soft_ids (optional): A tensor of weights (probabilities) used to
mix the embedding vectors.
kwargs: Additional keyword arguments for
:torch_nn:`functional.embedding` besides :attr:`params` and
:attr:`ids`.
Returns:
If :attr:`ids` is given, returns a Tensor of shape
``list(ids.shape) + embedding-dim``. For example,
if ``list(ids.shape) == [batch_size, max_time]``
and ``list(embedding.shape) == [vocab_size, emb_dim]``, then the
return tensor has shape ``[batch_size, max_time, emb_dim]``.
If :attr:`soft_ids` is given, returns a Tensor of shape
``list(soft_ids.shape)[:-1] + embedding-dim``. For example,
if ``list(soft_ids.shape) == [batch_size, max_time, vocab_size]``
and ``list(embedding.shape) == [vocab_size, emb_dim]``, then the
return tensor has shape ``[batch_size, max_time, emb_dim]``.
"""
if ids is not None:
if soft_ids is not None:
raise ValueError(
'Must not specify `ids` and `soft_ids` at the same time.')
ids_rank = ids.dim()
elif soft_ids is not None:
ids_rank = soft_ids.dim() - 1
else:
raise ValueError('Either `ids` or `soft_ids` must be given.')
embedding = self._embedding
if self._hparams.dropout_strategy == 'item_type':
noise_shape = self._get_noise_shape(self._hparams.dropout_strategy)
embedding = self._dropout_layer(embedding, noise_shape)
if ids is not None:
outputs = F.embedding(ids, embedding, **kwargs)
else:
outputs = embedder_utils.soft_embedding_lookup(embedding, soft_ids)
if self._hparams.dropout_strategy != 'item_type':
noise_shape = self._get_noise_shape(self._hparams.dropout_strategy,
ids_rank=ids_rank,
dropout_input=outputs)
outputs = self._dropout_layer(outputs, noise_shape)
return outputs
def __init__(self,
embedding,
start_tokens,
end_token,
tau,
embedding_size=None,
stop_gradient=False,
use_finish=True):
if callable(embedding):
self._embedding_fn = embedding
if embedding_size is None:
raise ValueError('`embedding_size` must be provided if '
'`embedding` is a callable.')
self._embedding_size = tf.convert_to_tensor(embedding_size,
dtype=tf.int32,
name="embedding_size")
else:
self._embedding_fn = (
lambda soft_ids: soft_embedding_lookup(embedding, soft_ids))
self._embedding_size = tf.shape(embedding)[0]
self._start_tokens = tf.convert_to_tensor(start_tokens,
dtype=tf.int32,
name="start_tokens")
self._end_token = tf.convert_to_tensor(end_token,
dtype=tf.int32,
name="end_token")
if self._start_tokens.get_shape().ndims != 1:
raise ValueError("start_tokens must be a vector")
self._batch_size = array_ops.size(start_tokens)
if self._end_token.get_shape().ndims != 0:
raise ValueError("end_token must be a scalar")
soft_start_tokens = tf.one_hot(self._start_tokens,
self._embedding_size,
dtype=tf.float32)
self._embedding_args_cnt = len(utils.get_args(self._embedding_fn))
if self._embedding_args_cnt == 1:
self._start_inputs = self._embedding_fn(soft_ids=soft_start_tokens)
elif self._embedding_args_cnt == 2:
# Position index is 0 in the beginning
times = tf.zeros([self._batch_size], dtype=tf.int32)
self._start_inputs = self._embedding_fn(soft_ids=soft_start_tokens,
times=times)
else:
raise ValueError('`embedding` should expect 1 or 2 arguments.')
self._batch_size = tf.size(self._start_tokens)
self._tau = tau
self._stop_gradient = stop_gradient
self._use_finish = use_finish
def _build(self,
ids=None,
soft_ids=None,
stop_gradient=False,
mode=None,
**kwargs):
"""Embeds (soft) ids.
Either :attr:`ids` or :attr:`soft_ids` must be given, and they
must not be given at the same time.
Args:
ids (optional): An integer tensor containing the ids to embed.
soft_ids (optional): A tensor of weights (probabilities) used to
mix the embedding vectors.
stop_gradient (bool): Whether to stop gradient back-propagation
to the embedding tensor. This can be used when, e.g., updating
only `soft_ids` while keeping the embedding tensor fixed.
Default to `False`.
mode (optional): A tensor taking value in
:tf_main:`tf.estimator.ModeKeys <estimator/ModeKeys>`, including
`TRAIN`, `EVAL`, and `PREDICT`. If `None`, dropout is
controlled by :func:`texar.global_mode`.
kwargs: Additional keyword arguments for
:tf_main:`tf.nn.embedding_lookup <nn/embedding_lookup>` besides
:attr:`params` and :attr:`ids`.
Returns:
If :attr:`ids` is given, returns a Tensor of shape
`shape(ids) + embedding-dim`. For example,
if `shape(ids) = [batch_size, max_time]`
and `shape(embedding) = [vocab_size, emb_dim]`, then the return
tensor has shape `[batch_size, max_time, emb_dim]`.
If :attr:`soft_ids` is given, returns a Tensor of shape
`shape(soft_ids)[:-1] + embdding-dim`. For example,
if `shape(soft_ids) = [batch_size, max_time, vocab_size]`
and `shape(embedding) = [vocab_size, emb_dim]`, then the return
tensor has shape `[batch_size, max_time, emb_dim]`.
"""
if ids is not None:
if soft_ids is not None:
raise ValueError(
'Must not specify `ids` and `soft_ids` at the same time.')
ids_rank = get_rank(ids)
elif soft_ids is not None:
ids_rank = get_rank(soft_ids) - 1
else:
raise ValueError('Either `ids` or `soft_ids` must be given.')
embedding = self._embedding
if stop_gradient:
embedding = tf.stop_gradient(embedding)
is_training = is_train_mode(mode)
if self._hparams.dropout_strategy == 'item_type':
dropout_layer = self._get_dropout_layer(self._hparams)
if dropout_layer:
embedding = dropout_layer.apply(inputs=embedding,
training=is_training)
if ids is not None:
outputs = tf.nn.embedding_lookup(embedding, ids, **kwargs)
else:
outputs = embedder_utils.soft_embedding_lookup(embedding, soft_ids)
if self._hparams.dropout_strategy != 'item_type':
dropout_layer = self._get_dropout_layer(self._hparams,
ids_rank=ids_rank,
dropout_input=outputs)
if dropout_layer:
outputs = dropout_layer.apply(inputs=outputs,
training=is_training)
return outputs