sampler: allow without replacement sampling for LogUniformSampler (#85)

Co-authored-by: Albert Zeyer <albzey@gmail.com>

Author: Gerstenberger

i6_models/parts/samplers/log_uniform.py

@@ -8,38 +8,43 @@
 
 class LogUniformSampler(nn.Module):
     def __init__(
-        self, num_classes: int, *, distribution_clamp_min: float = 1e-10, device: Optional[torch.device] = None
+        self,
+        num_classes: int,
+        *,
+        with_replacement: bool = True,
+        distribution_clamp_min: float = 1e-10,
+        device: Optional[torch.device] = None,
     ):
         """
-        Samples from a log uniform distribution from classes. Sampling is performed with replacement, i.e. sampled
-        indices can appear more than once in sampled set. This can be implemented with
-        `torch.distributions.multinomial.Multinomial` or `torch.multinomial`.
+        Samples from a log uniform distribution from classes. This assumes that the vocabulary is sorted according to
+        word count descending.
 
         :param num_classes: number of classes from which the distribution is sampled. The class indices are sorted in
             descending order according to their frequency.
+        :param with_replacement: wether to sample with replacement or not.
+        :param distribution_clamp_min: minimum probability mass.
         :param device: device on which the distribution is sampled.
         """
         super().__init__()
 
         # assumes count-sorted vocabulary, descending
         self.num_classes = num_classes
+        self.with_replacement = with_replacement
 
         # approximately zipf distribution
         ws = torch.arange(self.num_classes, dtype=torch.get_default_dtype(), device=device)
         self._distribution = (torch.log1p(ws + 1) - torch.log1p(ws)) / torch.log1p(torch.tensor(self.num_classes))
         self._distribution.clamp_(min=distribution_clamp_min)
         self._distribution /= self._distribution.sum()
 
-        self._cat_sampler = torch.distributions.categorical.Categorical(probs=self._distribution)
-
     def sample(self, num_samples: int) -> torch.Tensor:
         """
         Returns a random tensor in the size of [num_samples].
 
         :param num_samples: number of samples.
         :return: [num_samples]
         """
-        return self._cat_sampler.sample(torch.Size([num_samples]))
+        return torch.multinomial(self._distribution, num_samples, replacement=self.with_replacement)
 
     def log_prob(self, indices: torch.Tensor) -> torch.Tensor:
         """
@@ -48,4 +53,4 @@ def log_prob(self, indices: torch.Tensor) -> torch.Tensor:
         :param indices: the ground truth target labels as indices.
         :return: [B x T]
         """
-        return self._cat_sampler.log_prob(indices)
+        return torch.log(self._distribution[indices])