2-timeAwareNaive.py

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, CIFAR10
from torchvision.utils import save_image, make_grid
from tqdm import tqdm
import argparse
import os

class NoisePredictorTimeAwareNaive(nn.Module):
    def __init__(self, imageChannels, baseDepth=64):
        super().__init__()
        
        self.network = nn.Sequential(
            nn.Conv2d(imageChannels+1,  baseDepth,  3, padding=1), nn.BatchNorm2d(baseDepth),   nn.SiLU(), # +1 for time dim
            nn.Conv2d(baseDepth,      2*baseDepth,  3, padding=1), nn.BatchNorm2d(2*baseDepth), nn.SiLU(),
            nn.Conv2d(2*baseDepth,    4*baseDepth,  3, padding=1), nn.BatchNorm2d(4*baseDepth), nn.SiLU(),
            nn.Conv2d(4*baseDepth,    8*baseDepth,  3, padding=1), nn.BatchNorm2d(8*baseDepth), nn.SiLU(),
            nn.Conv2d(8*baseDepth,    4*baseDepth,  3, padding=1), nn.BatchNorm2d(4*baseDepth), nn.SiLU(),
            nn.Conv2d(4*baseDepth,    2*baseDepth,  3, padding=1), nn.BatchNorm2d(2*baseDepth), nn.SiLU(),
            nn.Conv2d(2*baseDepth,      baseDepth,  3, padding=1), nn.BatchNorm2d(baseDepth),   nn.SiLU(),
            nn.Conv2d(baseDepth,    imageChannels,  3, padding=1))

    def forward(self, x, t):
        timeChannel = t.view(-1, 1, 1, 1).expand(-1, 1, x.shape[-2], x.shape[-1])
        combinedInput = torch.cat((x, timeChannel), dim=1)
        return self.network(combinedInput)

class DiffusionModelTimeAwareNaive(nn.Module):
    def __init__(self, imageShape, betaStart=1e-4, betaEnd=2e-2, noisingSteps=1000, lr=1e-4, baseDepth=64):
        super().__init__()
        self.noisePredictor = NoisePredictorTimeAwareNaive(imageShape[0], baseDepth)
        self.imageShape     = imageShape
        self.noisingSteps   = noisingSteps

        self.optimizer = torch.optim.Adam(self.parameters(), lr=lr)
        self.level = 0

        beta        = torch.linspace(betaStart, betaEnd, noisingSteps) # beta[0] is noise variance at x_0 to x_1 transition
        alpha       = 1 - beta
        alphaBar    = torch.cumsum(torch.log(alpha), dim=0).exp()

        scheduleVariables = {
            "betaSqrt"                  : beta.sqrt(),
            "alpha"                     : alpha,
            "alphaBar"                  : alphaBar,
            "alphaBarSqrt"              : alphaBar.sqrt(),
            "alphaSqrtReciprocal"       : 1 / alpha.sqrt(),
            "sqrt1minusAlphaBar"        : (1 - alphaBar).sqrt(),
            "denoisingStepScaleFactor"  : (1 - alpha) / (1 - alphaBar).sqrt()}
        
        for key, value in scheduleVariables.items():
            self.register_buffer(key, value)

    @property
    def device(self):
        return next(self.parameters()).device # inferring device from where the class is stored

    def forward(self, nSamples):
        x_t = torch.randn(nSamples, *self.imageShape, dtype=torch.float32, device=self.device)  # x_T ~ N(0, I)
        for t in range(self.noisingSteps-1, -1, -1):
            noiseWeAdd = torch.randn(nSamples, *self.imageShape, device=self.device) if t > 0 else 0
            timestepEmbedding = torch.full((nSamples,), t/(self.noisingSteps-1)*2 - 1, device=self.device)
            totalNoisePredicted = self.noisePredictor(x_t, timestepEmbedding)
            x_t = self.alphaSqrtReciprocal[t]*(x_t - self.denoisingStepScaleFactor[t]*totalNoisePredicted) + self.betaSqrt[t]*noiseWeAdd
        return x_t

    def update(self, x_0):
        selectedTimesteps = torch.randint(0, self.noisingSteps, (x_0.shape[0],), device=self.device) # Downshifted by 1 to fit schedule indices
        noise = torch.randn_like(x_0)
        x_t = self.alphaBarSqrt[selectedTimesteps].view(-1, 1, 1, 1)*x_0 + self.sqrt1minusAlphaBar[selectedTimesteps].view(-1, 1, 1, 1)*noise
        timestepEmbedding = ((selectedTimesteps)/(self.noisingSteps-1))*2 - 1
        totalNoisePredicted = self.noisePredictor(x_t, timestepEmbedding)

        loss = F.mse_loss(totalNoisePredicted, noise)
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()
        return loss.item()
    
    def saveCheckpoint(self, path):
        data = {
            "model"     : self.state_dict(),
            "optimizer" : self.optimizer.state_dict(),
            "level"     : self.level}
        torch.save(data, path)

    def loadCheckpoint(self, path):
        if not path.endswith(".pth"): path += ".pth"
        data = torch.load(path, map_location=self.device)
        self.load_state_dict(data["model"])
        self.optimizer.load_state_dict(data["optimizer"])
        self.level = data["level"]

def ensurePath(*pathElements):
    path = os.path.join(*pathElements)
    os.makedirs(os.path.dirname(path) if os.path.splitext(path)[1] else path, exist_ok=True)
    return path

def main(args, device):
    if args.dataset == "cifar10":
        DATASET_CLASS = CIFAR10
    elif args.dataset == "mnist":
        DATASET_CLASS = MNIST

    print(f"Preparing dataset '{args.dataset}' (will download if missing)...")
    datasetTransform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))]) # from <0, 1> to <-1, 1>
    dataset = DATASET_CLASS(ensurePath("datasets"), download=True, transform=datasetTransform)
    imageShape = dataset[0][0].shape
    dataloader = DataLoader(dataset, batch_size=args.batchSize, shuffle=True)
    print(f"Starting run '{args.runName}': Training on {type(dataset).__name__} dataset with images of shape {tuple(imageShape)}")

    diffusionModel = DiffusionModelTimeAwareNaive(imageShape=imageShape, betaStart=1e-4, betaEnd=2e-2,
                                                  noisingSteps=args.noisingSteps,
                                                  lr=args.learningRate,
                                                  baseDepth=args.baseNetworkDepth).to(device)
    if args.checkpoint:
        diffusionModel.loadCheckpoint(ensurePath("checkpoints", args.checkpoint))
        print(f"Loaded the checkpoint {args.checkpoint}, so we're already at level {diffusionModel.level}")
        print(f"Loaded the checkpoint {args.checkpoint}, so we're already at level {diffusionModel.level}")

    lossSmooth = None
    for i in range(diffusionModel.level + 1, diffusionModel.level + args.epochs + 1):
        diffusionModel.train()
        progressBar = tqdm(dataloader)
        for x_0_batch, _ in progressBar: # labels are not needed
            loss = diffusionModel.update(x_0_batch.to(device))
            lossSmooth = 0.99*lossSmooth + 0.01*loss if lossSmooth is not None else loss
            progressBar.set_description(f"Epoch {i}, loss: {lossSmooth:.4f}")
        diffusionModel.level += 1

        if i % args.saveInterval == 0:
            if args.saveOutput:
                diffusionModel.eval()
                with torch.no_grad():
                    sampledImages = diffusionModel(16)
                    save_image(make_grid(sampledImages, nrow=4, normalize=True, value_range=(-1, 1)),
                            ensurePath("output", f"{args.runName}_level_{diffusionModel.level}.png"))
            if args.saveCheckpoints:
                checkpointsDir = ensurePath("checkpoints")
                currentPath  = os.path.join(checkpointsDir, f"{args.runName}_level_{diffusionModel.level}.pth")
                previousPath = os.path.join(checkpointsDir, f"{args.runName}_level_{diffusionModel.level-args.saveInterval}.pth")
                if os.path.exists(previousPath):
                    os.remove(previousPath)
                diffusionModel.saveCheckpoint(currentPath)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("-n",  "--runName",                 type=str,   default="coreDiffusionV2")
    parser.add_argument("-d",  "--dataset",                 type=str.lower, default="mnist", choices=["mnist", "cifar10"])
    parser.add_argument("-e",  "--epochs",                  type=int,   default=100)
    parser.add_argument("-o",  "--saveOutput",              action="store_true")
    parser.add_argument("-s",  "--saveCheckpoints",         action="store_true")
    parser.add_argument("-i",  "--saveInterval",            type=int,   default=1)
    parser.add_argument("-ch", "--checkpoint",              type=str,   default=None)
    parser.add_argument("-lr", "--learningRate",            type=float, default=1e-4)
    parser.add_argument("-b",  "--batchSize",               type=int,   default=64)
    parser.add_argument("-ns", "--noisingSteps",            type=int,   default=1000)
    parser.add_argument("-nd", "--baseNetworkDepth",        type=int,   default=64)

    device = "cuda" if torch.cuda.is_available() else "cpu"
    main(parser.parse_args(), device)