From 6562e19fc854c18c147a3246f91b497228b40a46 Mon Sep 17 00:00:00 2001 From: QB3 Date: Mon, 11 Dec 2023 16:11:16 -0500 Subject: [PATCH 1/2] modified generate samples --- examples/cifar10/utils_cifar.py | 30 +++++++++++++++++++----------- 1 file changed, 19 insertions(+), 11 deletions(-) diff --git a/examples/cifar10/utils_cifar.py b/examples/cifar10/utils_cifar.py index bc47cbb8..b60a8026 100644 --- a/examples/cifar10/utils_cifar.py +++ b/examples/cifar10/utils_cifar.py @@ -1,24 +1,32 @@ +import os import torch from torchdyn.core import NeuralODE +from tqdm import tqdm -# from torchvision.transforms import ToPILImage from torchvision.utils import make_grid, save_image use_cuda = torch.cuda.is_available() device = torch.device("cuda" if use_cuda else "cpu") -def generate_samples(node_, model, savedir, step, net_="normal"): +def generate_samples( + node_, model, savedir, step, net_="normal", batch_size=512, + num_gen=1000): model.eval() - with torch.no_grad(): - traj = node_.trajectory( - torch.randn(64, 3, 32, 32).to(device), - t_span=torch.linspace(0, 1, 100).to(device), - ) - traj = traj[-1, :].view([-1, 3, 32, 32]).clip(-1, 1) - traj = traj / 2 + 0.5 - save_image(traj, savedir + f"{net_}_generated_FM_images_step_{step}.png", nrow=8) - + gen_path = savedir + f"{net_}_generated_FM_images_step_{step}.png" + for batch in tqdm(range(num_gen//batch_size + 1)): + with torch.no_grad(): + traj = node_.trajectory( + torch.randn(64, 3, 32, 32).to(device), + t_span=torch.linspace(0, 1, 100).to(device), + ) + traj = traj[-1, :].view([-1, 3, 32, 32]).clip(-1, 1) + traj = traj / 2 + 0.5 + for i in range(batch_size): + idx = batch * batch_size + i + if idx < num_gen: + img_path = os.path.join(gen_path, f"{idx}.png") + save_image(traj[i], img_path) model.train() From c41f83da43d2e3826c832f65c2ead6f69bb1b413 Mon Sep 17 00:00:00 2001 From: QB3 Date: Tue, 12 Dec 2023 11:07:08 -0500 Subject: [PATCH 2/2] black formatting --- examples/cifar10/utils_cifar.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/examples/cifar10/utils_cifar.py b/examples/cifar10/utils_cifar.py index b60a8026..38986ff9 100644 --- a/examples/cifar10/utils_cifar.py +++ b/examples/cifar10/utils_cifar.py @@ -10,11 +10,11 @@ def generate_samples( - node_, model, savedir, step, net_="normal", batch_size=512, - num_gen=1000): + node_, model, savedir, step, net_="normal", batch_size=512, num_gen=1000 +): model.eval() gen_path = savedir + f"{net_}_generated_FM_images_step_{step}.png" - for batch in tqdm(range(num_gen//batch_size + 1)): + for batch in tqdm(range(num_gen // batch_size + 1)): with torch.no_grad(): traj = node_.trajectory( torch.randn(64, 3, 32, 32).to(device),