run.py

from tqdm import tqdm
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision.datasets import MNIST
from bern_models import ContextUnet, NaiveModel
from bernddpm import BernoulliDDPM
from dataset import MCDataset
import os


# hardcoding these here
n_epoch = 40
batch_size = 2048
lrate = 0.002  # 0.002
n_T = 10
betas=(0.05, 0.25)
n_feat = 16
img_size=8

device = torch.device("cpu")
save_dir = './saved_models/'

# model = NaiveModel(
#     N=img_size,
#     hidden_dim=n_feat
# )

model = ContextUnet(
    in_channels=1,
    img_size=img_size,
    n_feat=n_feat,
    cond_dim=2
)

ddpm = BernoulliDDPM(
    nn_model=model,
    betas=betas,
    n_T=n_T,
    device=device,
    image_size=img_size
)
ddpm.to(device)


def count_params(a, n):
    n_params = sum(p.numel() for p in a.parameters())
    print(f"{n:<20}: {n_params/1e3:.2f}K")

# count the number of parameters
n_params = sum(p.numel() for p in ddpm.parameters())
print(f"Total parameters: {n_params/1e3:.2f}K")


# optionally load a model
# ddpm.load_state_dict(torch.load("./data/diffusion_outputs/ddpm_unet01_mnist_9.pth"))

dataset = MCDataset(
    # data_base=data_base,
    num_images=10000,
    seed=23,
    N=img_size,
    Js=[-1, 1],
    # J=-1,
)
print(len(dataset))
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)


optim = torch.optim.Adam(ddpm.parameters(), lr=lrate)
# optim = torch.optim.SGD(ddpm.parameters(), lr=lrate)

for ep in range(n_epoch):
    ddpm.train()

    # linear lrate decay
    optim.param_groups[0]['lr'] = lrate*(1-ep/n_epoch)

    pbar = tqdm(dataloader)
    loss_ema = None
    for c, x in pbar:
        optim.zero_grad()
        x = x.to(device)
        c = c.to(device)
        loss = ddpm(x, c)
        loss.backward()
        loss_ema = 0.95 * loss_ema + 0.05 * loss.item() if loss_ema is not None else loss.item()
        pbar.set_description(f"E{ep+1:<2} loss: {loss_ema:.6f}")
        optim.step()

    # optionally save model
    if not (ep+1)%10:
        os.makedirs(save_dir, exist_ok=True)
        torch.save(ddpm.state_dict(), save_dir + f"model_{img_size}_{ddpm.nn_model.__class__.__name__}_{ep+1}.pth")
        print('saved model at ' + save_dir + f"model_{img_size}_{ddpm.nn_model.__class__.__name__}_{ep+1}.pth")