snsd0805/DDPM_Mnist
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix: set n in sample() default value to self.batch_size

Browse Source
This commit is contained in:
snsd0805 2023-03-14 22:57:44 +08:00
parent 37adcc7e97
commit a330666a17
Signed by: snsd0805
GPG Key ID: 569349933C77A854
1 changed files with 5 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
ddpm.py
View File

@ -53,14 +53,17 @@ class DDPM(nn.Module):
return mu + sigma * epsilon, epsilon # (b, c, w, h) return mu + sigma * epsilon, epsilon # (b, c, w, h)
def sample(self, model, generate_iteration_pic=False, n=self.batch_size): def sample(self, model, generate_iteration_pic=False, n=None):
''' '''
Inputs: Inputs:
model (nn.Module): Unet instance model (nn.Module): Unet instance
n (int): want to sample n pictures generate_iteration_pic (bool): whether generate 10 pic on different denoising time
n (int, default=self.batch_size): want to sample n pictures
Outputs: Outputs:
x_0 (nn.Tensor): (n, c, h, w) x_0 (nn.Tensor): (n, c, h, w)
''' '''
if n == None:
n = self.batch_size
c, h, w = 1, 28, 28 c, h, w = 1, 28, 28
model.eval() model.eval()
with torch.no_grad(): with torch.no_grad():