4-UNet.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 SinusoidalEmbeddings(nn.Module):
    def __init__(self, numEmbeddings, embeddingSize):
        super().__init__()
        position = torch.arange(numEmbeddings).unsqueeze(-1)
        divisor = torch.exp((torch.arange(0, embeddingSize, 2) / embeddingSize) * torch.log(torch.tensor(10000.0)))
        embeddings = torch.zeros(numEmbeddings, embeddingSize, requires_grad=False)
        embeddings[:, 0::2] = torch.sin(position / divisor)
        embeddings[:, 1::2] = torch.cos(position / divisor)
        self.register_buffer("embeddings", embeddings)

    def forward(self, t):
        return self.embeddings[t]

class ResBlock(nn.Module):
    def __init__(self, channelsIn, channelsOut, timestepEmbeddingSize):
        super().__init__()

        self.groupNorm1     = nn.GroupNorm(min(channelsIn//8, 16), channelsIn)
        self.activation1    = nn.SiLU()
        self.convolution1   = nn.Conv2d(channelsIn, channelsOut, kernel_size=3, padding=1)
        
        self.timestepPreprocessing = nn.Sequential(
            nn.SiLU(),
            nn.Linear(timestepEmbeddingSize, channelsOut * 2)) # AdaGN style, scale and shift for each channel
        
        self.groupNorm2     = nn.GroupNorm(min(channelsOut//8, 16), channelsOut)
        self.activation2    = nn.SiLU()
        self.convolution2   = nn.Conv2d(channelsOut, channelsOut, kernel_size=3, padding=1)

        self.residualConvolution = nn.Conv2d(channelsIn, channelsOut, kernel_size=1) if channelsIn != channelsOut else nn.Identity()

    def forward(self, x, timestepEmbedding):
        z = self.groupNorm1(x)
        z = self.activation1(z)
        z = self.convolution1(z)
        
        scale, shift = self.timestepPreprocessing(timestepEmbedding)[:, :, None, None].chunk(2, dim=1)
        
        z = self.groupNorm2(z)
        z = z*(1 + scale) + shift
        z = self.activation2(z)
        z = self.convolution2(z)

        return z + self.residualConvolution(x)

class NoisePredictorUNet(nn.Module):
    def __init__(self, imageChannels, timesteps, timestepEmbeddingSize=32, baseDepth=64):
        super().__init__()
        
        self.timestepEmbeddingsNet = nn.Sequential(
            SinusoidalEmbeddings(timesteps, timestepEmbeddingSize),
            nn.Linear(timestepEmbeddingSize, timestepEmbeddingSize),
            nn.SiLU(),
            nn.Linear(timestepEmbeddingSize, timestepEmbeddingSize))

        self.initialConvolution = nn.Conv2d(imageChannels, baseDepth, 3, padding=1)

        self.encoder1 = ResBlock(baseDepth, 2*baseDepth, timestepEmbeddingSize)
        self.downsampling1 = nn.Conv2d( 2*baseDepth, 2*baseDepth, 4, 2, 1) # MNIST 28x28 -> 14x14
        self.encoder2 = ResBlock(2*baseDepth, 4*baseDepth, timestepEmbeddingSize)
        self.downsampling2 = nn.Conv2d( 4*baseDepth, 4*baseDepth, 4, 2, 1) # MNIST 14x14 -> 7x7

        self.bottleneck1 = ResBlock(4*baseDepth, 8*baseDepth, timestepEmbeddingSize)
        self.bottleneck2 = ResBlock(8*baseDepth, 8*baseDepth, timestepEmbeddingSize)
        self.bottleneck3 = ResBlock(8*baseDepth, 4*baseDepth, timestepEmbeddingSize)

        self.upsampling1 = nn.ConvTranspose2d(4*baseDepth, 4*baseDepth, 4, 2, 1) # MNIST 7x7 -> 14x14
        self.decoder1 = ResBlock(4*baseDepth + 4*baseDepth, 2*baseDepth, timestepEmbeddingSize) # in + residual
        self.upsampling2 = nn.ConvTranspose2d(2*baseDepth, 2*baseDepth, 4, 2, 1) # MNIST 14x14 -> 28x28
        self.decoder2 = ResBlock(2*baseDepth + 2*baseDepth, baseDepth, timestepEmbeddingSize) # in + residual

        self.finalConvolution = nn.Conv2d(baseDepth, imageChannels, 3, padding=1)

    def forward(self, x, t):
        x = self.initialConvolution(x)
        timestepFeatures = self.timestepEmbeddingsNet(t)
        
        x = residual2 = self.encoder1(x, timestepFeatures)
        x = self.downsampling1(x)
        x = residual1 = self.encoder2(x, timestepFeatures)
        x = self.downsampling2(x)

        x = self.bottleneck1(x, timestepFeatures)
        x = self.bottleneck2(x, timestepFeatures)
        x = self.bottleneck3(x, timestepFeatures)

        x = self.upsampling1(x)
        x = self.decoder1(torch.cat((x, residual1), dim=1), timestepFeatures)

        x = self.upsampling2(x)
        x = self.decoder2(torch.cat((x, residual2), dim=1), timestepFeatures)

        return self.finalConvolution(x)

class DiffusionModelTimeAware(nn.Module):
    def __init__(self, imageShape, betaStart=1e-4, betaEnd=2e-2, noisingSteps=1000, lr=1e-4, timestepEmbeddingSize=32, baseDepth=64):
        super().__init__()
        self.noisePredictor = NoisePredictorUNet(imageShape[0], noisingSteps, timestepEmbeddingSize, 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, 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
            totalNoisePredicted = self.noisePredictor(x_t, torch.full((nSamples,), t, dtype=torch.int, device=self.device))
            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
        totalNoisePredicted = self.noisePredictor(x_t, selectedTimesteps)

        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 = DiffusionModelTimeAware(imageShape=imageShape, betaStart=1e-4, betaEnd=2e-2,
                                             noisingSteps=args.noisingSteps,
                                             lr=args.learningRate,
                                             timestepEmbeddingSize=args.timestepEmbeddingSize,
                                             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}")

    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="coreDiffusionV4")
    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("-te", "--timestepEmbeddingSize",   type=int,   default=128)
    parser.add_argument("-nd", "--baseNetworkDepth",        type=int,   default=64)

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