trajdl.loss.sampled_softmax module

class trajdl.loss.sampled_softmax.SampledSoftmaxLoss(weights: Tensor, bias: Tensor, num_words: int, num_samples: int, reduction: str = 'mean', device=None, use_sampled_softmax_in_eval: bool = False)

Bases: Module

Based on the default log_uniform_candidate_sampler in tensorflow. :param weights: :type weights: torch.Tensor :param bias: :type bias: torch.Tensor :param num_words: :type num_words: int :param num_samples: :type num_samples: int :param reduction: :type reduction: str :param device: :param use_sampled_softmax_in_eval: Whether or not using sampled softmax in eval mode, default False. :type use_sampled_softmax_in_eval: bool, optional

forward(embeddings: Tensor, targets: LongTensor, mask: BoolTensor = None) → Tensor

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

initialize_num_words(num_words: int)

log_uniform_candidate_sampler(targets: ~torch.LongTensor, choice_func=<function _choice>) → Tuple[LongTensor, Tensor, Tensor]

returns sampled, true_expected_count, sampled_expected_count

Parameters:

targets (shape is (batch_size,))

Returns:

• sampled_ids (shape is (n_samples,))

• true_expected_count (shape is (batch_size,))

• sampled_expected_count (shape is (n_samples,))