向前和向后扩散以实现稳定扩散

Question

我正在尝试测试前向扩散和后向扩散功能。我正在使用以下论文中描述的算法https://arxiv.org/abs/2006.11239

我的测试是拍摄图像。将其通过前向扩散过程。然后将其通过向后扩散过程。我期望获得原始图像。然而，去噪不起作用，我得到了随机图像。

我使用以下公式进行前向扩散：

和 ddpm 向后：

以下是我的代码测试：

import cv2
import numpy as np
MAX_NOISE_STEP = 20
# create a fixed beta schedule
beta = np.linspace(0.0001, 0.02, MAX_NOISE_STEP)

# this will be used as discussed in the reparameterization trick
alpha = 1 - beta
one_minus_alpha = 1-alpha
alpha_bar = np.cumprod(alpha, 0)
alpha_sqrt = np.sqrt(alpha)
alpha_bar = np.concatenate((np.array([1.]), alpha_bar[:-1]), axis=0)
sqrt_alpha_bar = np.sqrt(alpha_bar)
one_minus_alpha_bar = 1-alpha_bar
one_minus_sqrt_alpha_bar = np.sqrt(1-alpha_bar)



def stepNoise(key, x_t, t):
    np.random.seed(key)
    noise = np.random.normal(size=x_t.shape)
    betaT= np.take(beta,t)
    noisy_image =  ((1-betaT)**0.5) *x_t + betaT * noise
    return noisy_image,noise
# this function will add noise to the input as per the given timestamp
def forward_noise(key, x_0, t):
    np.random.seed(key)
    noise = np.random.normal(size=x_0.shape)
    noisy_image = np.take(sqrt_alpha_bar, t)  * x_0 + np.take(one_minus_alpha_bar, t)  * noise
    return noisy_image, noise
def ddpm(x_t, pred_noise, t):
    alpha_sqrt_t = np.take(alpha_sqrt, t)
    one_minus_alpha_t = np.take(one_minus_alpha, t)
    one_minus_sqrt_alpha_bar_t = np.take(one_minus_sqrt_alpha_bar, t)
    eps_coef = one_minus_alpha_t/ one_minus_sqrt_alpha_bar_t
    mean = (1 / alpha_sqrt_t) * (x_t - eps_coef * pred_noise)
    var = np.take(beta, t)
    z = np.random.normal(size=x_t.shape)

    return mean + (var ** .5) * z
    

#get feed
rng, tsrng = np.random.randint(0, 100000, size=(2,))
#read image and shape for batch
im = cv2.imread("10_29_2021.png")
im=np.reshape(im, (1,im.shape[0],im.shape[1],im.shape[2]))

im_c = ((im / 255.0) * 2 - 1)

noisyImage, noise = forward_noise(rng,im_c,MAX_NOISE_STEP-1)

x =noisyImage
for t in reversed(range(1,MAX_NOISE_STEP)):
    x= ddpm(x, noise, t)


x_image  = ((x + 1) / 2) * 255
noisyImage_image = ((noisyImage + 1) / 2) * 255

cv2.imshow('original',im[0])
cv2.imshow('Noisy',noisyImage[0])
cv2.imshow('denoised',x_image[0])
cv2.waitKey()

Answer 1

问题似乎是在后向扩散过程之前没有加载经过训练的模型，这对于预测每个时间步的噪声至关重要。如果未加载模型，反向扩散过程将无法正确对图像进行去噪，从而导致输出有噪声或不正确。 确保在开始逆向过程之前加载经过训练的模型。这很关键，因为模型需要预测向后扩散的每一步的噪声：

model = torch.load('path_to_trained_model.pth')
model.eval()  # Set the model to evaluation mode

在反向扩散循环中使用训练好的模型来预测噪声并对图像进行去噪：

for t in reversed(range(1, MAX_NOISE_STEP)):
    pred_noise = model(x, t)  # Predict the noise at timestep t 

x = ddpm(x, pred_noise, t) # 基于预测噪声去噪 如果不加载经过训练的模型，逆向过程将无法按预期运行。确保您正在加载模型并在反向过程中使用它。