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

feat: add sample.py

Browse Source
This commit is contained in:
snsd0805 2023-03-14 02:16:06 +08:00
parent 7a0b5a2a8a
commit 25e6a5ff62
Signed by: snsd0805
GPG Key ID: 569349933C77A854
1 changed files with 20 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
sample.py Normal file
View File

@ -0,0 +1,20 @@
import torch
import matplotlib.pyplot as plt
from ddpm import DDPM
from unet import Unet
BATCH_SIZE = 512
ITERATION = 1500
TIME_EMB_DIM = 128
DEVICE = torch.device('cuda')
model = Unet(TIME_EMB_DIM, DEVICE).to(DEVICE)
ddpm = DDPM(BATCH_SIZE, ITERATION, 1e-4, 2e-2, DEVICE)
model.load_state_dict(torch.load('unet.pth'))
x_t = ddpm.sample(model, 32)
for index, pic in enumerate(x_t):
p = pic.to('cpu').permute(1, 2, 0)
plt.imshow(p)
plt.savefig("output/{}.png".format(index))