diff options
| -rw-r--r-- | autoencoder.py | 11 | ||||
| -rw-r--r-- | main.py | 25 | ||||
| -rw-r--r-- | plots.py | 150 |
3 files changed, 101 insertions, 85 deletions
diff --git a/autoencoder.py b/autoencoder.py index f16fe79..645f628 100644 --- a/autoencoder.py +++ b/autoencoder.py @@ -11,6 +11,8 @@ from trainer import Trainer class PrintLayer(torch.nn.Module): def forward(self, x): print(x.shape) + print(x.min()) + print(x.max()) return x @@ -37,7 +39,7 @@ class Autoencoder(torch.nn.Module): torch.nn.ReLU(), torch.nn.Flatten(), torch.nn.Linear(64 * 7 * 7, self.latent_size), - torch.nn.ReLU(), + torch.nn.Sigmoid(), torch.nn.Unflatten(1, latent_dim), ) self.decoder = torch.nn.Sequential( @@ -55,7 +57,7 @@ class Autoencoder(torch.nn.Module): torch.nn.ReLU(), torch.nn.Conv2d(16, 16, kernel_size=3, padding=1), torch.nn.ReLU(), - torch.nn.Conv2d(16, input_dim, kernel_size=3, padding=1), + torch.nn.Conv2d(16, input_dim[0], kernel_size=3, padding=1), torch.nn.Sigmoid(), ) @@ -75,6 +77,7 @@ class AutoencoderDataset(Dataset): def __init__(self, dataset, device='cpu'): self.dataset = dataset self.device = device + self.dummy_param = torch.nn.Parameter(torch.empty(0)) def __len__(self): return len(self.dataset) @@ -100,7 +103,7 @@ def main(): # Initialize model model = Autoencoder(input_dim=(1, 28, 28), latent_dim=(1, 8, 8)) - # model.load_state_dict(torch.load('autoencoder.pth')) + model.load_state_dict(torch.load('autoencoder.pth')) model.to(device) # Train model @@ -108,7 +111,7 @@ def main(): lr = 1e-3 epochs = 10 optimizer = torch.optim.Adam(model.parameters(), lr=lr) - criterion = torch.nn.functional.mse_loss + criterion = torch.nn.functional.binary_cross_entropy trainer.train(model, dataloader, epochs, optimizer, criterion) # Save model @@ -185,8 +185,8 @@ class LatentDataset(Dataset): def __getitem__(self, index): img, label = self.dataset[index] - img = img.to(self.autoencoder.device) - latent = self.autoencoder.encode(img.unsqueeze(0)) + img = img.to(DEVICE) + latent = self.autoencoder.encode(img.unsqueeze(0)).squeeze(0) return latent, label def __len__(self): @@ -199,10 +199,10 @@ def loss(y_pred, y_true): def forward_diffusion(x0): x = x0.clone() - xs = [x.cpu()] + xs = [x] for t in range(1, DIFFU_STEPS+1): x = q_xt_xt_1(x, t)[0] - xs.append(x.cpu()) + xs.append(x) return xs @@ -238,14 +238,16 @@ dataset = datasets.MNIST( # dataset = FolderDataset('data/lfwcrop_color/faces') # dataset = FolderDataset('data/edface') -autoencoder = Autoencoder(1, 64).to(DEVICE) +autoencoder = Autoencoder((1, 28, 28), (1, 8, 8)).to(DEVICE) +autoencoder.load_state_dict(torch.load('autoencoder.pth')) +autoencoder.eval() dataset = LatentDataset(dataset, autoencoder) img = dataset[0][0] NB_CHANNEL, IMG_SIZE, _ = img.shape NB_LABEL = 10 -EPOCHS = 10 +EPOCHS = 100 LEARNING_RATE = 2e-4 @@ -287,7 +289,6 @@ if __name__ == '__main__': ############## with torch.no_grad(): - # Forward diffusion img, label = dataset[np.random.randint(0, len(dataset))] img = img.to(DEVICE) * 2 - 1 @@ -304,16 +305,16 @@ if __name__ == '__main__': plot_i = plots_id.index(t) plt.subplot(2, nb_plots + 1, plot_i + 2) plt.title(f"t={t}") - plt.imshow(tensor_to_image(x)) + plt.imshow(tensor_to_image(dataset.autoencoder.decode(x.unsqueeze(0)).squeeze(0))) plt.axis("off") for t in range(1, DIFFU_STEPS+1): - x = q_xt_x0(img, t)[0].cpu() + x = q_xt_x0(img, t)[0] if t not in plots_id: continue plot_i = plots_id.index(t) plt.subplot(2, nb_plots + 1, nb_plots + plot_i + 3) - plt.imshow(tensor_to_image(x)) + plt.imshow(tensor_to_image(dataset.autoencoder.decode(x.unsqueeze(0)).squeeze(0))) plt.axis("off") plt.subplot(2, nb_plots + 1, 1) @@ -348,7 +349,7 @@ if __name__ == '__main__': plt.subplot(n_classes, nb_plots, t_plot_i + nb_plots * class_i + 1) if class_i == 0: plt.title(f"t={t}") - plt.imshow(tensor_to_image(x[class_i])) + plt.imshow(tensor_to_image(dataset.autoencoder.decode(x[class_i].unsqueeze(0)).squeeze(0))) plt.axis("off") plt.tight_layout() plt.savefig("backward_diffusion.tmp.png") @@ -372,7 +373,7 @@ if __name__ == '__main__': for j in range(n_classes): id = i * n_classes + j plt.subplot(n_classes, nb_plots, id + 1) - plt.imshow(tensor_to_image(x[id])) + plt.imshow(tensor_to_image(dataset.autoencoder.decode(x[id].unsqueeze(0)).squeeze(0))) plt.axis("off") plt.tight_layout() plt.savefig("benchmark.tmp.png") @@ -7,6 +7,7 @@ import os from rich.progress import track from main import UNet, q_xt_xt_1, p_xt_1_xt, tensor_to_image +from autoencoder import Autoencoder os.makedirs('plots', exist_ok=True) @@ -27,6 +28,16 @@ BIN_MIN = -4 BIN_MAX = 4 + +model = UNet().to(DEVICE) +model.load_state_dict(torch.load('model.pth')) + +autoencoder = Autoencoder(input_dim=(1, 28, 28), latent_dim=(1, 8, 8)).to(DEVICE) +autoencoder.load_state_dict(torch.load('autoencoder.pth')) + + + + # plt.figure() # plt.plot(BETA, label='beta') # plt.plot(ALPHA, label='alpha') @@ -39,98 +50,99 @@ BIN_MAX = 4 mnist = datasets.MNIST('data', train=True, download=True) img, label = mnist[np.random.randint(0, len(mnist))] img = np.array(img) / 255 * 2 - 1 +img = torch.tensor(img, device=DEVICE, dtype=torch.float32).unsqueeze(0).unsqueeze(0) -# plt.figure() -# plt.imshow(img, cmap='gray') -# plt.title('Image') -# plt.axis('off') -# plt.savefig('plots/img.tmp.png') - +encoded = autoencoder.encode(img) -# plt.figure() -# plt.hist(img.flatten(), bins=NB_BINS, range=(BIN_MIN, BIN_MAX)) -# plt.yscale('log') -# plt.title('Image histogram') -# plt.savefig('plots/img_hist.tmp.png') +img = encoded.squeeze().cpu().detach().numpy() +plt.figure() +plt.imshow(img, cmap='gray') +plt.title('Image') +plt.axis('off') +plt.savefig('plots/img.tmp.png') -# def norm_dist(x, mean, std): -# return np.exp(-0.5 * ((x - mean) / std) ** 2) / (std * np.sqrt(2 * np.pi)) -# x_norm = np.linspace(BIN_MIN, BIN_MAX, 100) -# y_norm = norm_dist(x_norm, 0, 1) * 28**2 / NB_BINS * (BIN_MAX - BIN_MIN) +plt.figure() +plt.hist(img.flatten(), bins=NB_BINS, range=(BIN_MIN, BIN_MAX)) +plt.yscale('log') +plt.title('Image histogram') +plt.savefig('plots/img_hist.tmp.png') -# fig = plt.figure(figsize=(10, 5)) -# fig.suptitle('Diffusion naturelle') -# gs = GridSpec(1, 3, figure=fig) -# ax1 = fig.add_subplot(gs[0, 0]) -# ax2 = fig.add_subplot(gs[0, 1:]) +def norm_dist(x, mean, std): + return np.exp(-0.5 * ((x - mean) / std) ** 2) / (std * np.sqrt(2 * np.pi)) -# plots_to_save = np.linspace(1, DIFFU_STEPS, 100).astype(int) +x_norm = np.linspace(BIN_MIN, BIN_MAX, 100) +y_norm = norm_dist(x_norm, 0, 1) * 8**2 / NB_BINS * (BIN_MAX - BIN_MIN) -# xt = torch.tensor(img, device=DEVICE, dtype=torch.float32).unsqueeze(0).unsqueeze(0) -# for t in track(range(1, DIFFU_STEPS+1)): -# xt, eps = q_xt_xt_1(xt, t) +fig = plt.figure(figsize=(10, 5)) +fig.suptitle('Diffusion naturelle') +gs = GridSpec(1, 3, figure=fig) +ax1 = fig.add_subplot(gs[0, 0]) +ax2 = fig.add_subplot(gs[0, 1:]) -# if t not in plots_to_save: -# continue +plots_to_save = np.linspace(1, DIFFU_STEPS, 100).astype(int) -# xt_numpy = xt.cpu().detach().numpy()[0, 0] +xt = torch.tensor(img, device=DEVICE, dtype=torch.float32).unsqueeze(0).unsqueeze(0) +for t in track(range(1, DIFFU_STEPS+1)): + xt, eps = q_xt_xt_1(xt, t) -# ax1.clear() -# ax1.imshow(xt_numpy, cmap='gray') -# ax1.set_title(f'xt at t={t:04d}') -# ax1.axis('off') + if t not in plots_to_save: + continue -# ax2.clear() -# ax2.hist(xt_numpy.flatten(), bins=NB_BINS, range=(BIN_MIN, BIN_MAX)) -# ax2.plot(x_norm, y_norm, color='red', label='N(0, 1)') -# ax2.legend() -# ax2.set_yscale('log') -# ax2.set_ylim(y_norm.min(), 1e3) -# ax2.set_title(f'xt histogram') + xt_numpy = xt.cpu().detach().numpy()[0, 0] -# fig.savefig(f'plots/diffusion/{t:04d}.tmp.png') -# os.system('convert -delay 20 -loop 0 plots/diffusion/*.png plots/diffusion.tmp.gif') + ax1.clear() + ax1.imshow(xt_numpy, cmap='gray') + ax1.set_title(f'xt at t={t:04d}') + ax1.axis('off') + ax2.clear() + ax2.hist(xt_numpy.flatten(), bins=NB_BINS, range=(BIN_MIN, BIN_MAX)) + ax2.plot(x_norm, y_norm, color='red', label='N(0, 1)') + ax2.legend() + ax2.set_yscale('log') + ax2.set_ylim(y_norm.min(), 1e3) + ax2.set_title(f'xt histogram') + fig.savefig(f'plots/diffusion/{t:04d}.tmp.png') +os.system('convert -delay 20 -loop 0 plots/diffusion/*.png plots/diffusion.tmp.gif') -model = UNet().to(DEVICE) -model.load_state_dict(torch.load('mnist_model.pth')) +fig = plt.figure(figsize=(10, 5)) +fig.suptitle('Diffusion inverse') +gs = GridSpec(1, 3, figure=fig) +ax1 = fig.add_subplot(gs[0, 0]) +ax2 = fig.add_subplot(gs[0, 1:]) -# fig = plt.figure(figsize=(10, 5)) -# fig.suptitle('Diffusion inverse') -# gs = GridSpec(1, 3, figure=fig) -# ax1 = fig.add_subplot(gs[0, 0]) -# ax2 = fig.add_subplot(gs[0, 1:]) +xt = torch.randn(1, 1, 28, 28, device=DEVICE) +vec = torch.zeros(1, 10).to(DEVICE) +vec[0, label] = 1 +for t in track(range(DIFFU_STEPS, 0, -1)): + t_tensor = torch.tensor([[t]], device=DEVICE, dtype=torch.float32) + xt = p_xt_1_xt(model, xt, t_tensor, vec) -# xt = torch.randn(1, 1, 28, 28, device=DEVICE) -# vec = torch.zeros(1, 10).to(DEVICE) -# vec[0, label] = 1 -# for t in track(range(DIFFU_STEPS, 0, -1)): -# t_tensor = torch.tensor([[t]], device=DEVICE, dtype=torch.float32) -# xt = p_xt_1_xt(model, xt, t_tensor, vec) + if t not in plots_to_save: + continue -# if t not in plots_to_save: -# continue + xt_numpy = xt.cpu().detach().numpy()[0, 0] -# xt_numpy = xt.cpu().detach().numpy()[0, 0] + ax1.clear() + ax1.imshow(xt_numpy, cmap='gray') + ax1.set_title(f'xt at t={t:04d}') + ax1.axis('off') -# ax1.clear() -# ax1.imshow(xt_numpy, cmap='gray') -# ax1.set_title(f'xt at t={t:04d}') -# ax1.axis('off') + ax2.clear() + ax2.hist(xt_numpy.flatten(), bins=NB_BINS, range=(BIN_MIN, BIN_MAX)) + ax2.plot(x_norm, y_norm, color='red', label='N(0, 1)') + ax2.legend() + ax2.set_yscale('log') + ax2.set_ylim(y_norm.min(), 1e3) + ax2.set_title(f'xt histogram') -# ax2.clear() -# ax2.hist(xt_numpy.flatten(), bins=NB_BINS, range=(BIN_MIN, BIN_MAX)) -# ax2.plot(x_norm, y_norm, color='red', label='N(0, 1)') -# ax2.legend() -# ax2.set_yscale('log') -# ax2.set_ylim(y_norm.min(), 1e3) -# ax2.set_title(f'xt histogram') + fig.savefig(f'plots/diffusion_inverse/{t:04d}.tmp.png') +os.system('convert -delay 20 -loop 0 -reverse plots/diffusion_inverse/*.png plots/diffusion_inverse.tmp.gif') -# fig.savefig(f'plots/diffusion_inverse/{t:04d}.tmp.png') -# os.system('convert -delay 20 -loop 0 -reverse plots/diffusion_inverse/*.png plots/diffusion_inverse.tmp.gif') +exit(0) tpause = {150: 'xt', 20: 'mu', 50: 'xt_1'} |