Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[Feature] Support uniform timesteps sampler for DDPM #153

Merged
merged 2 commits into from
Nov 29, 2021
Merged

[Feature] Support uniform timesteps sampler for DDPM #153

Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
463849c
support uniform timesteps sampler for DDPM
LeoXing1996 Oct 31, 2021
61516cd
fix known issuses
LeoXing1996 Nov 12, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
fix known issuses
  • Loading branch information
@LeoXing1996
LeoXing1996 committed Nov 12, 2021
commit 61516cd27ca2fd337ba35b9765f0120e529f7d55
8 changes: 6 additions & 2 deletions mmgen/models/diffusions/sampler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ class UniformTimeStepSampler:

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.
Expand All @@ -24,9 +25,12 @@ def sample(self, batch_size):
Returns:
torch.Tensor: Sampled timesteps.
"""
p = [1 / self.num_timesteps for _ in range(self.num_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=p))
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)