add ddpm utface data load/train example

Author: AkiRusProd

.gitignore

@@ -136,3 +136,5 @@ dmypy.json
 *.npy
 *.png
 *.jpeg
+*.zip
+*.nt

README.md

@@ -145,6 +145,18 @@ Some [examples](examples/) were trained on the [MNIST](https://pjreddie.com/proj
 
 #### More details about some of them:
 
+<details>
+<summary>Denoising Diffusion Probabilistic Model (DDPM)</summary>
+
+<p align="center">
+<img src="generated images/ddpm_mnist.gif" width=20% height=20%>
+<img src="generated images/ddpm_utkface.gif" width=20% height=20%>
+</p>
+
+Code:   
+*[Model Example](examples/ddpm.py)*
+</details>
+
 <details>
 <summary>Convolutional Classifier</summary>
 

data_loader.py

@@ -1,3 +1,4 @@
+import zipfile
 from pathlib import Path
 
 import numpy as np
@@ -6,13 +7,13 @@
 from mnist_data_downloader import download_data
 
 
-def prepare_data(data):
+def prepare_mnist_data(data):
     inputs, targets = [], []
 
     for raw_line in tqdm(data, desc="preparing data"):
         line = raw_line.split(",")
 
-        inputs.append(np.asfarray(line[1:]))
+        inputs.append(np.asfarray(line[1:]) / 127.5 - 1)  # normalization: / 255 => [0; 1]  #/ 127.5-1 => [-1; 1]
         targets.append(int(line[0]))
 
     return inputs, targets
@@ -31,10 +32,10 @@ def load_mnist(path="datasets/mnist/"):
 
 
     if not (Path(path) / "mnist_train.npy").exists() or not (Path(path) / "mnist_test.npy").exists():
-        training_inputs, training_targets = prepare_data(training_data)
+        training_inputs, training_targets = prepare_mnist_data(training_data)
         training_inputs = np.asfarray(training_inputs)
 
-        test_inputs, test_targets = prepare_data(test_data)
+        test_inputs, test_targets = prepare_mnist_data(test_data)
         test_inputs = np.asfarray(test_inputs)
 
         np.save(path + "mnist_train.npy", training_inputs)
@@ -53,3 +54,47 @@ def load_mnist(path="datasets/mnist/"):
     test_dataset = test_inputs
 
     return training_dataset, test_dataset, training_targets, test_targets
+
+
+import os
+
+
+def prepare_utkface_data(path, image_size = (3, 32, 32)):
+        
+    import random
+
+    import numpy as np
+    from PIL import Image
+    
+    images = os.listdir(path)
+    random.shuffle(images)
+    
+    training_inputs = []
+    for image in tqdm(images, desc = 'preparing data'):
+        image = Image.open(path + "/" + image)
+        image = image.resize((image_size[1], image_size[2]))
+        image = np.asarray(image)
+        image = image.transpose(2, 0, 1)
+        image = image / 127.5 - 1
+        training_inputs.append(image)
+
+    return np.array(training_inputs)
+
+
+def load_utkface(path="datasets/utkface/", image_size=(3, 32, 32)):
+    path = Path(path)
+    if not path.exists():
+        path.mkdir(parents=True)
+
+    if not (path / 'UTKFace').exists():
+        with zipfile.ZipFile(path / 'archive.zip', 'r') as zip_ref:
+            zip_ref.extractall(path)
+
+    save_path = path / 'UTKFace.npy'
+    if not save_path.exists():
+        training_inputs = prepare_utkface_data(path / 'UTKFace', image_size)
+        np.save(save_path, training_inputs)
+    else:
+        training_inputs = np.load(save_path)
+
+    return training_inputs

datasets/utkface/readme.txt

@@ -0,0 +1 @@
+For training examples you need to download the zip from the link: https://www.kaggle.com/datasets/jangedoo/utkface-new and drop it into this folder.

examples/ddpm.py

@@ -12,7 +12,7 @@
 
 import neunet as nnet
 import neunet.nn as nn
-from data_loader import load_mnist
+from data_loader import load_mnist, load_utkface  # noqa F401
 from neunet import Tensor
 from neunet.optim import Adam
 
@@ -274,7 +274,7 @@ def denormalize(x):
         else:
             return Image.fromarray(images_array)
 
-    def train(self, dataset, epochs, batch_size, image_path, image_size, save_every_epochs=1):
+    def train(self, dataset, epochs, batch_size, image_path, image_size, save_path, save_every_epochs=1):
         channels, H_size, W_size = image_size
 
         data_batches = np.array_split(dataset, np.arange(batch_size, len(dataset), batch_size))
@@ -333,6 +333,11 @@ def train(self, dataset, epochs, batch_size, image_path, image_size, save_every_
                     loop=0,
                 )
 
+                if not Path(save_path).exists():
+                    Path(save_path).mkdir(parents=True, exist_ok=True)
+
+                nnet.save(self.model.state_dict(), f"{save_path}/ddpm_{epoch + 1}.nt")
+
             loss_history.append(epoch_loss)
 
         return loss_history
@@ -506,29 +511,32 @@ def forward(self, x, t):
 
 
 device = "cuda"
+image_size = (3, 32, 32)
+# image_size = (1, 28, 28) # for mnist
 
+training_data = load_utkface(image_size=(3, 32, 32))
+# training_data, _, _, _ = load_mnist() # for mnist
 
 diffusion = Diffusion(
     model=SimpleUNet(
-        image_channels=1,
-        image_size=28,
-        down_channels=(32, 64, 128),
-        up_channels=(128, 64, 32),
+        image_channels=image_size[0],
+        image_size = image_size[2],
+        down_channels=(128, 256, 512, 1024), # (32, 64, 128) for mnist
+        up_channels=(1024, 512, 256, 128), # (128, 64, 32) for mnist
     ).to(device),
     timesteps=300,
     beta_start=0.0001,
     beta_end=0.02,
     criterion=nn.MSELoss(),
 )
 
-training_data, test_data, training_labels, test_labels = load_mnist()
-training_data = training_data / 127.5 - 1  # normalization: / 255 => [0; 1]  #/ 127.5-1 => [-1; 1]
+# diffusion.model.load_state_dict(nnet.load("saved models/utkface_model/ddpm_3.nt")) # load saved model example if it exists
 
-# diffusion.ddpm_denoise_sample(25, (1, 28, 28))
 diffusion.train(
     training_data,
     epochs=3,
-    batch_size=16,
+    batch_size=5,
     image_path="generated images",
-    image_size=(1, 28, 28),
+    save_path = "saved models/utkface_model", # for "saved models/mnist_model" mnist
+    image_size=image_size,
 )