1. Description
Cet article est un tutoriel sur la création de votre propre modèle de diffusion à partir de zéro. J'aime toujours garder les choses simples et faciles, c'est pourquoi nous avons évité ici les mathématiques compliquées. Ce n’est pas un modèle de diffusion normal. Au lieu de cela, je l’appelle le modèle de diffusion rapide. Seuls les réseaux de neurones convolutifs (CNN) seront utilisés pour réaliser le modèle de diffusion. Dans cet article, je ne vous fournirai aucun fichier de modèle/poids/script existant.
Vous devez former le modèle vous-même.
(Nous utilisons l'ensemble de données CIFAR-10 fourni par TensorFlow.
Vous pouvez trouver le code dans mon GitHub
https://github.com/Seachaos/Tree.Rocks/blob/main/QuickDiffusionModel/QuickDiffusionModel.ipynb
2. Cette idée
Voici comment fonctionne le modèle de diffusion : c'est comme prendre une image complètement bruitée et améliorer progressivement la qualité de l'image jusqu'à ce qu'elle devienne claire.
(Comme indiqué ci-dessous)
Par conséquent, nous pouvons créer un modèle d'apprentissage en profondeur qui peut améliorer la qualité de l'image (du bruit complet aux images claires).
Pour une compréhension plus claire, veuillez consulter cet organigramme ci-joint.
Comme indiqué ci-dessus, le modèle tente de générer des images avec progressivement moins de bruit. Il ne nous reste plus qu’à former un modèle d’apprentissage profond pour apprendre à réduire le bruit.
Pour cette tâche, nous avons besoin de deux entrées du modèle :
- Image d'entrée - l'image bruyante doit être traitée
- horodatage - indique au modèle quel est l'état du bruit afin qu'il puisse l'apprendre plus facilement
3. Mettre en œuvre un modèle de diffusion rapide
Tout d’abord, importons ce dont nous avons besoin :
import numpy as np
from tqdm.auto import trange, tqdm
import matplotlib.pyplot as plt
import tensorflow as tf
from tensorflow.keras import layers Et préparez notre ensemble de données. Dans ce didacticiel, nous utiliserons une grande collection d'images de voitures (CIFAR-10) comme exemple pour garder les choses aussi simples et rapides que possible.
(Cependant, si vous disposez de suffisamment d’échantillons, vous pouvez choisir n’importe quelle image de votre choix.
(X_train, y_train), (X_test, y_test) = tf.keras.datasets.cifar10.load_data()
X_train = X_train[y_train.squeeze() == 1]
X_train = (X_train / 127.5) - 1.0Ensuite, définissons les variables.
IMG_SIZE = 32 # input image size, CIFAR-10 is 32x32
BATCH_SIZE = 128 # for training batch size
timesteps = 16 # how many steps for a noisy image into clear
time_bar = 1 - np.linspace(0, 1.0, timesteps + 1) # linspace for timestepsIci, nous définissons le « pas de temps », ce qui signifie que notre modèle apprendra à générer des images du bruit (niveau 0) au clair (niveau 16) tout au long du processus de formation.
Voyons une photo pour avoir une idée plus claire
plt.plot(time_bar, label='Noise')
plt.plot(1 - time_bar, label='Clarity')
plt.legend()
Comme vous pouvez le constater, du pas de temps 0 à 16, le bruit diminue et la clarté augmente progressivement. C'est ce que nous voulons que notre modèle apprenne.
et préparer quelques fonctions pour prévisualiser les données
def cvtImg(img):
img = img - img.min()
img = (img / img.max())
return img.astype(np.float32)
def show_examples(x):
plt.figure(figsize=(10, 10))
for i in range(25):
plt.subplot(5, 5, i+1)
img = cvtImg(x[i])
plt.imshow(img)
plt.axis('off')
show_examples(X_train)
CIFAR-10 Automobile
3.1 Préparation à la formation
Ici, nous devons préparer le code pour entraîner les images.
L'idée est d'obtenir deux images (A et B) à partir de points aléatoires dans le temps, où A est l'image bruyante et B est l'image la plus claire.
Notre modèle apprendra à convertir A en B (de bruyant à plus clair) en fonction de ce moment spécifique.
(encore une fois comme cette photo)
Nous avons donc ici la fonction forward_noise .
def forward_noise(x, t):
a = time_bar[t] # base on t
b = time_bar[t + 1] # image for t + 1
noise = np.random.normal(size=x.shape) # noise mask
a = a.reshape((-1, 1, 1, 1))
b = b.reshape((-1, 1, 1, 1))
img_a = x * (1 - a) + noise * a
img_b = x * (1 - b) + noise * b
return img_a, img_b
def generate_ts(num):
return np.random.randint(0, timesteps, size=num)
# t = np.full((25,), timesteps - 1) # if you want see clarity
# t = np.full((25,), 0) # if you want see noisy
t = generate_ts(25) # random for training data
a, b = forward_noise(X_train[:25], t)
show_examples(a)Si vous voulez comprendre comment cela fonctionne, je vous recommande d'exécuter le code que j'ai commenté. ( t = ... )
3.2 Construire des blocs CNN
Nous utiliserons U-Net comme modèle, les détails sont expliqués dans le code ci-dessous.
Architecture du modèle, les détails seront expliqués dans le code suivant. Avant de construire le modèle, nous devons d'abord définir les blocs.
Voici le code du bloc make :
def block(x_img, x_ts):
x_parameter = layers.Conv2D(128, kernel_size=3, padding='same')(x_img)
x_parameter = layers.Activation('relu')(x_parameter)
time_parameter = layers.Dense(128)(x_ts)
time_parameter = layers.Activation('relu')(time_parameter)
time_parameter = layers.Reshape((1, 1, 128))(time_parameter)
x_parameter = x_parameter * time_parameter
# -----
x_out = layers.Conv2D(128, kernel_size=3, padding='same')(x_img)
x_out = x_out + x_parameter
x_out = layers.LayerNormalization()(x_out)
x_out = layers.Activation('relu')(x_out)
return x_out Chaque bloc contient deux réseaux convolutifs avec des paramètres temporels, permettant au réseau de déterminer son pas de temps actuel et de générer les informations correspondantes.
Vous pouvez voir l'organigramme :
(x_img est l'image d'entrée, est l'image du bruit, x_ts est l'entrée pour le pas de temps)
Construire le modèle Maintenant, nous pouvons construire notre modèle
def make_model():
x = x_input = layers.Input(shape=(IMG_SIZE, IMG_SIZE, 3), name='x_input')
x_ts = x_ts_input = layers.Input(shape=(1,), name='x_ts_input')
x_ts = layers.Dense(192)(x_ts)
x_ts = layers.LayerNormalization()(x_ts)
x_ts = layers.Activation('relu')(x_ts)
# ----- left ( down ) -----
x = x32 = block(x, x_ts)
x = layers.MaxPool2D(2)(x)
x = x16 = block(x, x_ts)
x = layers.MaxPool2D(2)(x)
x = x8 = block(x, x_ts)
x = layers.MaxPool2D(2)(x)
x = x4 = block(x, x_ts)
# ----- MLP -----
x = layers.Flatten()(x)
x = layers.Concatenate()([x, x_ts])
x = layers.Dense(128)(x)
x = layers.LayerNormalization()(x)
x = layers.Activation('relu')(x)
x = layers.Dense(4 * 4 * 32)(x)
x = layers.LayerNormalization()(x)
x = layers.Activation('relu')(x)
x = layers.Reshape((4, 4, 32))(x)
# ----- right ( up ) -----
x = layers.Concatenate()([x, x4])
x = block(x, x_ts)
x = layers.UpSampling2D(2)(x)
x = layers.Concatenate()([x, x8])
x = block(x, x_ts)
x = layers.UpSampling2D(2)(x)
x = layers.Concatenate()([x, x16])
x = block(x, x_ts)
x = layers.UpSampling2D(2)(x)
x = layers.Concatenate()([x, x32])
x = block(x, x_ts)
# ----- output -----
x = layers.Conv2D(3, kernel_size=1, padding='same')(x)
model = tf.keras.models.Model([x_input, x_ts_input], x)
return model
model = make_model()
# model.summary()Il s'agit d'un U-Net. Pour les parties gauche, droite et MLP, veuillez vous référer à la figure ci-dessus (architecture du modèle).
N'oubliez pas de compiler le modèle
optimizer = tf.keras.optimizers.Adam(learning_rate=0.0008)
loss_func = tf.keras.losses.MeanAbsoluteError()
model.compile(loss=loss_func, optimizer=optimizer)Nous utilisons Adam comme optimiseur et MeanAbsoluteError (MAE) comme fonction de perte.
Résultats de la prédiction : nous pouvons maintenant essayer notre première prédiction. Les étapes de prédiction sont les suivantes :
- Créer des images bruyantes
- entrée dans notre modèle par pas de temps
- Continuez ainsi jusqu'à la fin du pas de temps
Voici donc la fonction :
def predict(x_idx=None):
x = np.random.normal(size=(32, IMG_SIZE, IMG_SIZE, 3))
for i in trange(timesteps):
t = i
x = model.predict([x, np.full((32), t)], verbose=0)
show_examples(x)
predict()
Image de sortie du modèle non entraîné Ci-dessus se trouve notre sortie de modèle non entraîné, et comme vous pouvez le voir, elle ne fait rien d'utile. Cette fonction nous aide également à visualiser chaque étape :
def predict_step():
xs = []
x = np.random.normal(size=(8, IMG_SIZE, IMG_SIZE, 3))
for i in trange(timesteps):
t = i
x = model.predict([x, np.full((8), t)], verbose=0)
if i % 2 == 0:
xs.append(x[0])
plt.figure(figsize=(20, 2))
for i in range(len(xs)):
plt.subplot(1, len(xs), i+1)
plt.imshow(cvtImg(xs[i]))
plt.title(f'{i}')
plt.axis('off')
predict_step()
4. Modèle de formation
Cette fonction de formation est très simple
def train_one(x_img):
x_ts = generate_ts(len(x_img))
x_a, x_b = forward_noise(x_img, x_ts)
loss = model.train_on_batch([x_a, x_ts], x_b)
return lossNous avons seulement besoin de fournir x_ts et x_img (x_a) pour que notre modèle puisse apprendre à générer x_b.
et en faire une fonction d'époque
def train(R=50):
bar = trange(R)
total = 100
for i in bar:
for j in range(total):
x_img = X_train[np.random.randint(len(X_train), size=BATCH_SIZE)]
loss = train_one(x_img)
pg = (j / total) * 100
if j % 5 == 0:
bar.set_description(f'loss: {loss:.5f}, p: {pg:.2f}%')Enfin, exécutez plusieurs fois et réduisez progressivement le taux d'apprentissage
for _ in range(10):
train()
# reduce learning rate for next training
model.optimizer.learning_rate = max(0.000001, model.optimizer.learning_rate * 0.9)
# show result
predict()
predict_step()
plt.show()Vous pouvez obtenir une image de sortie comme celle-ci
5. Conclusion
Ce tutoriel est conçu pour être simple et vous permettre d'expérimenter. Vous pouvez essayer vos propres paramètres (comme la modification de la taille de l'image, les filtres CNN, les pas de temps ou MLP, etc.) et vous entraîner avec plus d'époques pour obtenir de meilleurs résultats. Chaos de la mer