[Feature] Support uniform timesteps sampler for DDPM (#153)

* support uniform timesteps sampler for DDPM

* fix known issuses

mmgen/models/diffusions/__init__.py

@@ -0,0 +1,3 @@
+from .sampler import UniformTimeStepSampler
+
+__all__ = ['UniformTimeStepSampler']

mmgen/models/diffusions/sampler.py

@@ -0,0 +1,36 @@
+import numpy as np
+import torch
+
+from ..builder import MODULES
+
+
+@MODULES.register_module()
+class UniformTimeStepSampler:
+    """Timestep sampler for DDPM-based models. This sampler sample all
+    timesteps with the same probabilistic.
+
+    Args:
+        num_timesteps (int): Total timesteps of the diffusion process.
+    """
+
+    def __init__(self, num_timesteps):
+        self.num_timesteps = num_timesteps
+        self.prob = [1 / self.num_timesteps for _ in range(self.num_timesteps)]
+
+    def sample(self, batch_size):
+        """Sample timesteps.
+        Args:
+            batch_size (int): The desired batch size of the sampled timesteps.
+
+        Returns:
+            torch.Tensor: Sampled timesteps.
+        """
+        # use numpy to make sure our implementation is consistent with the
+        # official ones.
+        return torch.from_numpy(
+            np.random.choice(
+                self.num_timesteps, size=(batch_size, ), p=self.prob))
+
+    def __call__(self, batch_size):
+        """Return sampled results."""
+        return self.sample(batch_size)

tests/test_models/test_base_ddpm.py

@@ -0,0 +1,18 @@
+import torch
+
+from mmgen.models.diffusions import UniformTimeStepSampler
+
+
+def test_uniform_sampler():
+    sampler = UniformTimeStepSampler(10)
+    timesteps = sampler(2)
+    assert timesteps.shape == torch.Size([
+        2,
+    ])
+    assert timesteps.max() < 10 and timesteps.min() >= 0
+
+    timesteps = sampler.__call__(2)
+    assert timesteps.shape == torch.Size([
+        2,
+    ])
+    assert timesteps.max() < 10 and timesteps.min() >= 0