mnist_train_step_t.py

from denoising_diffusion_pytorch_step_t import Unet, GaussianDiffusion, Trainer
import torchvision
import os
import errno
import shutil

def create_folder(path):
    try:
        os.mkdir(path)
    except OSError as exc:
        if exc.errno != errno.EEXIST:
            raise
        pass

def del_folder(path):
    try:
        shutil.rmtree(path)
    except OSError as exc:
        pass

create = 0

if create:
    trainset = torchvision.datasets.MNIST(
            root='./data', train=True, download=True)
    root = './root_mnist/'
    del_folder(root)
    create_folder(root)

    for i in range(10):
        lable_root = root + str(i) + '/'
        create_folder(lable_root)

    for idx in range(len(trainset)):
        img, label = trainset[idx]
        img.save(root + str(label) + '/' + str(idx) + '.png')

timesteps=1000

model = Unet(
    dim = 64,
    dim_mults = (1, 2, 4, 8),
    image_size = 32,
    timesteps = timesteps,
    with_time_emb = False
).cuda()


diffusion = GaussianDiffusion(
    model,
    image_size = 32,
    timesteps = timesteps,   # number of steps
    loss_type = 'l1'    # L1 or L2
).cuda()

trainer = Trainer(
    diffusion,
    './root_mnist/',
    image_size = 32,
    train_batch_size = 32,
    train_lr = 2e-5,
    train_num_steps = 700000,         # total training steps
    gradient_accumulate_every = 2,    # gradient accumulation steps
    ema_decay = 0.995,                # exponential moving average decay
    fp16 = True,                       # turn on mixed precision training with apex
    results_folder = './results_mnist_step_t'
)

trainer.train()