pl_module_rectifiedflow.py

import torch
import math

import data_utils
import torch.utils.data

import imageio.v3 as imageio
import lightning.pytorch as pl
import torch.nn as nn
import torch.distributions as dist
import numpy as np
import safetensors.torch as st

from network_diffusion_unet import ConditionalUNet, ConditionalUNetDiT
from loss_fn import L1andGDL
from adam_atan2_pytorch import AdamAtan2
from lightning.pytorch.loggers.tensorboard import TensorBoardLogger
from lightning.pytorch.utilities import grad_norm
from lightning.pytorch.callbacks import LearningRateMonitor, StochasticWeightAveraging, LearningRateFinder
from torchvision.utils import make_grid

def convert_uniform_to_custom(u):
    #return 0.5 - torch.cos((1/3) * torch.acos(1 - 2 * u) + math.pi / 3)
    return 0.5 + 2 * torch.cos((2 * math.pi - torch.arccos((11/16)*(1-2*u)))/3)

class PLModule(pl.LightningModule):
    def __init__(self, mid_visual_ridge, mid_visual_basins, mid_visual_gt):
        super().__init__()
        self.save_hyperparameters()
        self.lr = 6e-4
        self.wd = 5e-5
        self.model = ConditionalUNetDiT(base_ch=8, embd_dim=16)
        #self.map_average = torch.from_numpy(imageio.imread(map_average)).unsqueeze(0)
        #self.map_average = (self.map_average - self.map_average.mean()) / self.map_average.std()
        self.loss_fn = L1andGDL()
        self.val_metrics = []
        self.mid_visual_ridge, self.mid_visual_basins = mid_visual_ridge, mid_visual_basins
        self.mid_visual_gt = mid_visual_gt
        self.initialize_model()

    def initialize_model(self):
        for name, m in self.model.named_modules():
            if isinstance(m, nn.Linear) and ('time_affine' in name or 'water_level_affine' in name):
                m.weight.data.zero_()
                m.bias.data.zero_()

    def configure_optimizers(self):
        opt = AdamAtan2(self.parameters(), lr=self.lr, decoupled_wd=True, weight_decay=self.wd)
        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(opt, 100, eta_min=1e-7)
        return {
            "optimizer": opt,
            "lr_scheduler": {"scheduler": scheduler, "interval": "epoch", "frequency": 1},
        }

    def _step(self, batch, batch_idx):
        x0, ridge_map, basin_map, water_level = batch
        b = water_level.shape[0]
        #map_average = self.map_average.expand((b, -1, -1, -1)).to(self.device)

        noise = torch.randn_like(x0, device=self.device, dtype=x0.dtype)
        t = torch.rand((b,), device=self.device)
        t = convert_uniform_to_custom(t).to(x0.dtype)

        xt = t.view(-1, 1, 1, 1) * x0 + (1 - t.view(-1, 1, 1, 1)) * noise
        v = x0 - noise

        predicted_v = self.model(xt, ridge_map, basin_map, water_level, t)  # Predict velocity v
        loss = self.loss_fn(predicted_v, v)  # Loss between predicted and target v
        return loss

    def training_step(self, batch, batch_idx):
        loss = self._step(batch, batch_idx)
        self.logger.experiment.add_scalar(f"Train/Loss", loss.detach(), self.global_step)
        return loss

    def validation_step(self, batch, batch_idx):
        loss = self._step(batch, batch_idx)
        self.val_metrics.append(loss.detach())
        return loss

    @torch.no_grad()
    def inference_step(self, ridge_map, basin_map, water_level, num_steps=50):
        device = self.device
        b = ridge_map.shape[0]
        x = torch.randn_like(ridge_map, device=device)
        water_level = torch.tensor((water_level,), device=device).expand(b,)
        time = torch.linspace(0, 1, num_steps + 1, device=device)

        for i in range(num_steps):
            t = torch.full((b,), time[i], device=device)
            dt = torch.full((b, 1, 1, 1), time[i+1] - time[i], device=device)

            v = self.model(x, ridge_map, basin_map, water_level, t)

            x = x + dt * v

        return x

    def on_train_epoch_end(self):
        sea_level = 0.0
        ridge_map = torch.from_numpy(imageio.imread(self.mid_visual_ridge))[None,None,:].to(device=self.device, dtype=torch.float32)

        basin_map = torch.from_numpy(imageio.imread(self.mid_visual_basins))[None,None,:].to(device=self.device)
        basin_map = (basin_map>=sea_level).to(torch.float32)
        output = self.inference_step(ridge_map, basin_map, sea_level)
        mid_visual_result = output.squeeze([1])
        self.logger.experiment.add_scalar("Visualize/Min", mid_visual_result.min(), self.current_epoch)
        self.logger.experiment.add_scalar("Visualize/Max", mid_visual_result.max(), self.current_epoch)
        self.logger.experiment.add_scalar("Visualize/Mean", mid_visual_result.mean(), self.current_epoch)
        mid_visual_result = (mid_visual_result - mid_visual_result.min()) / (mid_visual_result.max() - mid_visual_result.min())
        self.logger.experiment.add_image(f'Visualize/Model Output', mid_visual_result, self.current_epoch)

        vram_data = torch.cuda.mem_get_info()
        vram_usage = (vram_data[1] - vram_data[0]) / (1024 ** 2)
        self.logger.experiment.add_scalar(f"Other/VRAM Usage", vram_usage, self.current_epoch)
        torch.cuda.reset_peak_memory_stats()
        if self.current_epoch == 0:
            mid_visual_gt = torch.from_numpy(imageio.imread(self.mid_visual_gt))[None,:]
            mid_visual_gt = (mid_visual_gt - mid_visual_gt.min()) / (mid_visual_gt.max() - mid_visual_gt.min())
            self.logger.experiment.add_image(f'Visualize/Ridge', ridge_map.squeeze([1]), self.current_epoch)
            self.logger.experiment.add_image(f'Visualize/Basin', basin_map.squeeze([1]), self.current_epoch)
            self.logger.experiment.add_image(f'Visualize/GT', mid_visual_gt, self.current_epoch)

    def on_validation_epoch_end(self):
        epoch_averages = torch.stack(self.val_metrics).nanmean(dim=0)
        self.logger.experiment.add_scalar("Val/Loss", epoch_averages, self.current_epoch)
        self.val_metrics.clear()
    #def on_before_optimizer_step(self, optimizer):
    #    norms = grad_norm(self.model, norm_type=2)
    #    self.log_dict(norms, logger=True)



# Example usage
if __name__ == "__main__":
    torch.set_float32_matmul_precision('medium')
    if torch.cuda.is_available() and torch.version.cuda:
        print('Optimising computing and memory use via cuDNN! (NVIDIA GPU only).')
        torch.backends.cudnn.enabled = True
        torch.backends.cudnn.benchmark = True
        torch.backends.cudnn.allow_tf32 = True
    elif torch.cuda.is_available() and torch.version.hip:
        print('Optimising computing using TunableOp! (AMD GPU only).')
        torch.cuda.tunable.enable()
        torch.cuda.tunable.set_filename('TunableOp_results')

    train_split, val_split = data_utils.make_dataset_t_v('dataset_large')

    callbacks = []
    callbacks.append(LearningRateMonitor(logging_interval='epoch'))
    model_checkpoint = pl.callbacks.ModelCheckpoint(dirpath="", filename="FlashScape",
                                                    save_weights_only=False,
                                                    enable_version_counter=False, save_last=False)
    callbacks.append(model_checkpoint)
    swa_callback = StochasticWeightAveraging(1e-5, 0.8, int(0.2 * 100 - 1))
    callbacks.append(swa_callback)
    #lr_finder = LearningRateFinder(1e-5, 0.1)
    #callbacks.append(lr_finder)
    #model = PLModule.load_from_checkpoint('FlashScape V2.ckpt')
    trainer = pl.Trainer(max_epochs=100, log_every_n_steps=1, logger=TensorBoardLogger(f'lightning_logs', name='FlashScape Dit No MapAvg Zero Init'),
                         accelerator="gpu", enable_checkpointing=True,
                         precision='16-mixed', enable_progress_bar=True, num_sanity_val_steps=0, callbacks=callbacks)
    with trainer.init_module():
        model = PLModule('dataset_large/Ridge_11417648.tiff',
                         'dataset_large/Basins_11417648.tiff',
                         'dataset_large/11417648.tiff')
    model = torch.compile(model)


    train_dataset = data_utils.TrainDataset(train_split)
    val_dataset = data_utils.ValDataset(val_split)

    train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=8,
                                               num_workers=8, pin_memory=False, persistent_workers=True, shuffle=True)
    val_loader = torch.utils.data.DataLoader(dataset=val_dataset, batch_size=8,
                                             num_workers=8, pin_memory=False, persistent_workers=True)

    trainer.fit(model,
                val_dataloaders=val_loader,
                train_dataloaders=train_loader)
                #ckpt_path='FlashScape.ckpt')