Digit Generation by DDPM

ddpm.py

import torch

from torch.utils.data import Dataset
from torchvision import datasets, transforms
import torch.optim as optim

import cv2
from tqdm import tqdm

from denoising_diffusion_pytorch import Unet, GaussianDiffusion
import numpy as np

class Mydataset(Dataset):
    def __init__(self, mnist_data):
        self.mnist = mnist_data

    def __len__(self):
        return len(self.mnist)

    def __getitem__(self, idx):
        X = self.mnist[idx][0]
        X = (X * 2) - 1
        y = self.mnist[idx][1]

        return X, y

mnist_data = datasets.MNIST('.',
                            transform = transforms.Compose([
                                transforms.Resize(32,transforms.InterpolationMode.NEAREST),
                                transforms.ToTensor()
                                ]),
                            download=True)

mnist_data_norm = Mydataset(mnist_data)
dataloader = torch.utils.data.DataLoader(mnist_data_norm,
                                          batch_size=512,
                                          shuffle=True)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

model = Unet(
    dim = 16,
    dim_mults = (1, 2, 4, 8),
    channels=1
).to(device)
diffusion = GaussianDiffusion(
    model,
    image_size = 32,
    timesteps = 1000,   # number of steps
    loss_type = 'l1',
    channels=1
).to(device)
optimizer = optim.Adam(diffusion.parameters())

for epoch in range(2000):
    for i, data in enumerate(tqdm(dataloader), 0):
        optimizer.zero_grad()
        batch = data[0].to(device)
        loss = diffusion(batch)
        loss.backward()
        optimizer.step()

    imgs = diffusion.sample(batch_size = 16)
    imgs = imgs.cpu()
    imgs_img = np.empty((4 * imgs.shape[2], 4 * imgs.shape[3]), dtype=np.float32)
    for y in range(4):
        for x in range(4):
            imgs_img[y * 32:(y + 1) * 32, x * 32:(x + 1) * 32] = imgs[y*4+x]

    imgs_img = np.array(imgs_img * 255, dtype=np.uint8)
    cv2.imwrite(f"./MNIST/ddpm_{epoch}.bmp", imgs_img)