Delete pl_module_rectifiedflow.py

pl_module_rectifiedflow.py

@@ -1,181 +0,0 @@
-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')