diff options
| -rw-r--r-- | main.py | 36 | ||||
| -rw-r--r-- | test.py | 95 |
2 files changed, 9 insertions, 122 deletions
@@ -78,16 +78,11 @@ class UNet(nn.Module): DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") -DIFFU_STEPS = 20 -BETA = torch.linspace(1e-2, 3e-1, DIFFU_STEPS, device=DEVICE) +DIFFU_STEPS = 300 +BETA = torch.linspace(1e-4, 2e-2, DIFFU_STEPS, device=DEVICE) BETA = torch.cat((torch.tensor([0.], device=DEVICE), BETA)) ALPHA = 1 - BETA ALPHA_BAR = torch.cumprod(ALPHA, dim=0) -CUM_SQ_SUM = [] -for t in range(len(ALPHA)): - cum_sq_sum = sum([torch.prod(ALPHA[s+2:t+1]) * (1 - ALPHA[s+1])**2 for s in range(t)]) - CUM_SQ_SUM.append(cum_sq_sum) -CUM_SQ_SUM = torch.tensor(CUM_SQ_SUM, device=DEVICE) def q_xt_xt_1(xt_1, t): @@ -95,29 +90,17 @@ def q_xt_xt_1(xt_1, t): alpha = ALPHA[t_ind] mean = torch.sqrt(alpha) * xt_1 - std = 1 - alpha + std = torch.sqrt(1 - alpha) xt = torch.distributions.Normal(mean, std).sample() return xt -# def q_xt_x0(x0, t): -# t_ind = t.to(dtype=torch.long) if isinstance(t, torch.Tensor) else t - -# alpha = ALPHA[t_ind] -# alpha_bar = ALPHA_BAR[t_ind] -# alpha_bar_bar = ALPHA_BAR_BAR[t_ind - 1] -# mean = torch.sqrt(alpha_bar) * x0 -# std = alpha_bar_bar * (1 - alpha) -# return torch.distributions.Normal(mean, std).sample() - def q_xt_x0(x0, t): t_ind = t.to(dtype=torch.long) if isinstance(t, torch.Tensor) else t alpha_bar = ALPHA_BAR[t_ind] - cum_sq_sum = CUM_SQ_SUM[t_ind] - mean = torch.sqrt(alpha_bar) * x0 - std = torch.sqrt(cum_sq_sum) + std = torch.sqrt(1 - alpha_bar) return torch.distributions.Normal(mean, std).sample() @@ -156,8 +139,8 @@ class MNISTDiffusionDataset(Dataset): img, label = self.mnist_data[index] img = img.to(DEVICE) - # Normalize the image between -0.5 and 0.5. - img = img - 0.5 + # Normalize the image. + img = img * 2 - 1 # Add noise to the image. t = torch.randint(1, DIFFU_STEPS, (1,), device=DEVICE) @@ -204,9 +187,8 @@ if __name__ == '__main__': transform=transforms.ToTensor(), ) - img, label = mnist_data[1] - img = img.to(DEVICE) - 0.5 + img = img.to(DEVICE) * 2 - 1 nb_plots = 10 plots_id = [i for i in np.linspace(1, DIFFU_STEPS, nb_plots, dtype=int)] @@ -244,7 +226,7 @@ if __name__ == '__main__': model.load_state_dict(torch.load('model.pth')) # Define the optimizer. - optimizer = torch.optim.Adam(model.parameters(), lr=1e-4) + optimizer = torch.optim.Adam(model.parameters(), lr=2e-4) # Define the trainiself.encodevec = nn.Linear(10, 28*28)ng dataset. train_dataset = MNISTDiffusionDataset(train=True) @@ -252,7 +234,7 @@ if __name__ == '__main__': num_workers=4, persistent_workers=True) trainer = Trainer() criterion = loss - epochs = 1 + epochs = 10 # Train the model. trainer.train(model, train_loader, epochs, optimizer, criterion) diff --git a/test.py b/test.py deleted file mode 100644 index d1db16d..0000000 --- a/test.py +++ /dev/null @@ -1,95 +0,0 @@ -import numpy as np -import matplotlib.pyplot as plt - -A = np.random.uniform(0.1, 0.9) -B = np.random.uniform(0.1, 0.9) - -def q_xt_xt_1_simple(x, t): - mean = A * x - std = B - return np.random.normal(mean, std) - -def q_xt_x0_simple_damien(x, t): - mean = A ** t * x - std = np.sqrt(sum([A ** (2*i) for i in range(t)])) * B - return np.random.normal(mean, std) - - -def q_xt_xt_1(x, t): - alpha = ALPHA[t] - mean = np.sqrt(alpha) * x - std = 1 - alpha - return np.random.normal(mean, std) - -def q_xt_x0_paper(x, t): - alpha_bar = ALPHA_BAR[t] - mean = np.sqrt(alpha_bar) * x - std = 1 - alpha_bar - return np.random.normal(mean, std) - -def q_xt_x0_damien(x, t): - alpha_bar = ALPHA_BAR[t] - cum_sq_sum = sum([np.prod(ALPHA[s+2:t+1]) * (1 - ALPHA[s+1])**2 for s in range(t)]) - mean = np.sqrt(alpha_bar) * x - std = np.sqrt(cum_sq_sum) - return np.random.normal(mean, std) - -T = 100 -BETA = np.concatenate(([0], np.linspace(1e-4, 2e-2, T))) -ALPHA = 1 - BETA -ALPHA_BAR = np.cumprod(ALPHA) - -N = int(1e6) -x0 = 1 - -xs_implicit = np.array([x0] * N) -for t in range(1, T+1): - xs_implicit = q_xt_xt_1(xs_implicit, t) - -xs_explicit_paper = q_xt_x0_paper(np.array([x0] * N), T) -xs_explicit_damien = q_xt_x0_damien(np.array([x0] * N), T) - -plt.figure() -bins = np.linspace(min( - xs_implicit.min(), - xs_explicit_paper.min(), - xs_explicit_damien.min(), - ), max( - xs_implicit.max(), - xs_explicit_paper.max(), - xs_explicit_damien.max(), - ), 100) -plt.hist(xs_implicit, bins=bins, alpha=0.5, label="q_xt_xt_1") -plt.hist(xs_explicit_paper, bins=bins, alpha=0.5, label="q_xt_x0_paper") -plt.hist(xs_explicit_damien, bins=bins, alpha=0.5, label="q_xt_x0_damien") -plt.legend() -plt.savefig("q_xt_xt_1_vs_q_xt_x0.tmp.png") -plt.show() - -xs_implicit = np.array([x0] * N) -for t in range(1, T+1): - xs_implicit = q_xt_xt_1_simple(xs_implicit, t) - -xs_explicit_damien = q_xt_x0_simple_damien(np.array([x0] * N), T) - -plt.figure() -bins = np.linspace(min(xs_implicit.min(), xs_explicit_damien.min()), max(xs_implicit.max(), xs_explicit_damien.max()), 100) -plt.hist(xs_implicit, bins=bins, alpha=0.5, label="q_xt_xt_1_simple") -plt.hist(xs_explicit_damien, bins=bins, alpha=0.5, label="q_xt_x0_simple_damien") -plt.legend() -plt.savefig("q_xt_xt_1_simple_vs_q_xt_x0_simple.tmp.png") -plt.show() - -A1, B1, A2, B2, A3, B3 = np.random.uniform(0, 1, 6) -x1 = np.random.normal(A1, B1, N) -x2 = np.random.normal(A2 * x1, B2, N) -x3 = np.random.normal(A3 * x2, B3, N) -x3_ = np.random.normal(A1 * A2 * A3, np.sqrt(A3**2 * A2**2 * B1**2 + A3**2 * B2**2 + B3**2), N) - -plt.figure() -bins = np.linspace(min(x3.min(), x3_.min()), max(x3.max(), x3_.max()), 100) -plt.hist(x3, bins=bins, alpha=0.5, label="normal") -plt.hist(x3_, bins=bins, alpha=0.5, label="product") -plt.legend() -plt.savefig("product_normal.tmp.png") -plt.show() |