9-smoothWeights.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
from contextlib import contextmanager

class trailingModel(nn.Module):
    def __init__(self, model, smoothingFactor=0.999):
        super().__init__()
        self.smoothingFactor = smoothingFactor

        for name, parameter in model.named_parameters():
            self.register_buffer(name.replace('.', '_'), parameter.data.clone())

    @torch.no_grad()
    def update(self, model):
        for name, parameter in model.named_parameters():
            trailingParameter = getattr(self, name.replace('.', '_'))
            trailingParameter.lerp_(parameter.data, 1 - self.smoothingFactor)

    @contextmanager
    def swap(self, model):
        originalWeights = [p.data.clone() for p in model.parameters()]
        for name, parameter in model.named_parameters():
            parameter.data.copy_(getattr(self, name.replace('.', '_')))
        try:
            yield
        finally:
            for parameter, originalParameter in zip(model.parameters(), originalWeights):
                parameter.data.copy_(originalParameter)

@torch.no_grad()
def zeroInitialization(network):
    for parameter in network.parameters():
        parameter.zero_()
    return network

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(),
            zeroInitialization(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   = zeroInitialization(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 SelfAttention(nn.Module):
    def __init__(self, channels):
        super().__init__()
        self.layerNorm = nn.LayerNorm(channels)
        self.multiheadAttention = nn.MultiheadAttention(embed_dim=channels, num_heads=4, batch_first=True)
        zeroInitialization(self.multiheadAttention.out_proj)

    def forward(self, x):
        B, C, H, W = x.shape

        xFlat = x.flatten(start_dim=2, end_dim=-1).transpose(1, 2) # (B, HW, C)
        xNorm = self.layerNorm(xFlat)
        xOut = xFlat + self.multiheadAttention(xNorm, xNorm, xNorm, need_weights=False)[0]
        return xOut.transpose(1, 2).view(B, C, H, W)

class NoisePredictorAttentionUNet(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.bottleneckAttention = SelfAttention(8*baseDepth)
        self.bottleneck2 = ResBlock(8*baseDepth, 8*baseDepth, timestepEmbeddingSize)
        self.bottleneck3 = ResBlock(8*baseDepth, 4*baseDepth, timestepEmbeddingSize)

        self.upsampling1 = nn.Sequential(
            nn.Upsample(scale_factor=2, mode='nearest'),
            nn.Conv2d(4*baseDepth, 4*baseDepth, 3, padding=1)) # MNIST 7x7 -> 14x14
        self.decoder1 = ResBlock(4*baseDepth + 4*baseDepth, 2*baseDepth, timestepEmbeddingSize) # in + residual
        self.decoderAttention = SelfAttention(2*baseDepth)
        self.upsampling2 = nn.Sequential(
            nn.Upsample(scale_factor=2, mode='nearest'),
            nn.Conv2d(2*baseDepth, 2*baseDepth, 3, padding=1)) # MNIST 14x14 -> 28x28
        self.decoder2 = ResBlock(2*baseDepth + 2*baseDepth, baseDepth, timestepEmbeddingSize) # in + residual

        self.finalConvolution = zeroInitialization(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.bottleneckAttention(x)
        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.decoderAttention(x)
        x = self.upsampling2(x)
        x = self.decoder2(torch.cat((x, residual2), dim=1), timestepFeatures)

        return self.finalConvolution(x)

class DiffusionModelCosineSchedule(nn.Module):
    def __init__(self, imageShape, noisingSteps=1000, lr=1e-4, timestepEmbeddingSize=32, smoothingFactor=0.999, baseDepth=64):
        super().__init__()
        self.noisePredictor = NoisePredictorAttentionUNet(imageShape[0], noisingSteps, timestepEmbeddingSize, baseDepth)
        self.imageShape     = imageShape
        self.noisingSteps   = noisingSteps

        self.optimizer = torch.optim.Adam(self.parameters(), lr=lr)
        self.trailingNoisePredictor = trailingModel(self.noisePredictor, smoothingFactor=smoothingFactor)
        self.level = 0

        allSteps = torch.linspace(0, noisingSteps, noisingSteps + 1)
        f_t = torch.cos(((allSteps/noisingSteps + 0.008)/1.008) * torch.pi/2) ** 2
        alphaBar = f_t/f_t[0]

        beta = torch.clamp(1 - alphaBar[1:]/alphaBar[:-1], max=0.999)
        alpha = 1 - beta
        alphaBar = alphaBar[1:] # We want index 0 to be the first noising step, not original img
        
        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):
        with self.trailingNoisePredictor.swap(self.noisePredictor):
            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()
        self.trailingNoisePredictor.update(self.noisePredictor)
        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 = DiffusionModelCosineSchedule(imageShape=imageShape,
                                                  noisingSteps=args.noisingSteps,
                                                  lr=args.learningRate,
                                                  timestepEmbeddingSize=args.timestepEmbeddingSize,
                                                  smoothingFactor=args.smoothingFactor,
                                                  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="coreDiffusionV9")
    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("-sf", "--smoothingFactor",         type=float, default=0.999)
    parser.add_argument("-nd", "--baseNetworkDepth",        type=int,   default=64)

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