from omegaconf import OmegaConf import os from shutil import copyfile import warnings from typing import Dict,Sequence,Union import inspect import numpy as np import torch from torch.optim.lr_scheduler import LambdaLR, CosineAnnealingLR, SequentialLR import torchmetrics from lightning.pytorch import Trainer, loggers as pl_loggers from lightning.pytorch.profilers import PyTorchProfiler from lightning.pytorch.strategies import DDPStrategy from lightning.pytorch.callbacks import ( Callback, LearningRateMonitor, DeviceStatsMonitor, EarlyStopping, ModelCheckpoint ) from lightning.pytorch.utilities import grad_norm from einops import rearrange from models.vae import AutoencoderKL from models.knowledge_alignment import SEVIRAvgIntensityAlignment,get_alignment_kwargs_avg_x from models.diffusion import LatentDiffusion from models.core_model.cuboid_transformer import CuboidTransformerUNet from datamodule import SEVIRLightningDataModule,vis_sevir_seq from utils.path import ( default_exps_dir, default_pretrained_vae_dir,default_pretrained_alignment_dir ) from utils.optim import disable_train,warmup_lambda from utils.layout import step_layout_to_in_out_slice from evaluation import FrechetVideoDistance,SEVIRSkillScore class PreDiffSEVIRPLModule(LatentDiffusion): def __init__(self, total_num_steps: int, oc_file: str = None, save_dir: str = None): self.total_num_steps = total_num_steps if oc_file is not None: oc_from_file = OmegaConf.load(open(oc_file, "r")) else: oc_from_file = None oc = self.get_base_config(oc_from_file=oc_from_file) self.save_hyperparameters(oc) self.oc = oc latent_model_cfg = OmegaConf.to_object(oc.model.latent_model) num_blocks = len(latent_model_cfg["depth"]) if isinstance(latent_model_cfg["self_pattern"], str): block_attn_patterns = [latent_model_cfg["self_pattern"]] * num_blocks else: block_attn_patterns = OmegaConf.to_container(latent_model_cfg["self_pattern"]) latent_model = CuboidTransformerUNet( input_shape=latent_model_cfg["input_shape"], target_shape=latent_model_cfg["target_shape"], base_units=latent_model_cfg["base_units"], scale_alpha=latent_model_cfg["scale_alpha"], num_heads=latent_model_cfg["num_heads"], attn_drop=latent_model_cfg["attn_drop"], proj_drop=latent_model_cfg["proj_drop"], ffn_drop=latent_model_cfg["ffn_drop"], # inter-attn downsample/upsample downsample=latent_model_cfg["downsample"], downsample_type=latent_model_cfg["downsample_type"], upsample_type=latent_model_cfg["upsample_type"], upsample_kernel_size=latent_model_cfg["upsample_kernel_size"], # attention depth=latent_model_cfg["depth"], block_attn_patterns=block_attn_patterns, # global vectors num_global_vectors=latent_model_cfg["num_global_vectors"], use_global_vector_ffn=latent_model_cfg["use_global_vector_ffn"], use_global_self_attn=latent_model_cfg["use_global_self_attn"], separate_global_qkv=latent_model_cfg["separate_global_qkv"], global_dim_ratio=latent_model_cfg["global_dim_ratio"], # misc ffn_activation=latent_model_cfg["ffn_activation"], gated_ffn=latent_model_cfg["gated_ffn"], norm_layer=latent_model_cfg["norm_layer"], padding_type=latent_model_cfg["padding_type"], checkpoint_level=latent_model_cfg["checkpoint_level"], pos_embed_type=latent_model_cfg["pos_embed_type"], use_relative_pos=latent_model_cfg["use_relative_pos"], self_attn_use_final_proj=latent_model_cfg["self_attn_use_final_proj"], # initialization attn_linear_init_mode=latent_model_cfg["attn_linear_init_mode"], ffn_linear_init_mode=latent_model_cfg["ffn_linear_init_mode"], ffn2_linear_init_mode=latent_model_cfg["ffn2_linear_init_mode"], attn_proj_linear_init_mode=latent_model_cfg["attn_proj_linear_init_mode"], conv_init_mode=latent_model_cfg["conv_init_mode"], down_linear_init_mode=latent_model_cfg["down_up_linear_init_mode"], up_linear_init_mode=latent_model_cfg["down_up_linear_init_mode"], global_proj_linear_init_mode=latent_model_cfg["global_proj_linear_init_mode"], norm_init_mode=latent_model_cfg["norm_init_mode"], # timestep embedding for diffusion time_embed_channels_mult=latent_model_cfg["time_embed_channels_mult"], time_embed_use_scale_shift_norm=latent_model_cfg["time_embed_use_scale_shift_norm"], time_embed_dropout=latent_model_cfg["time_embed_dropout"], unet_res_connect=latent_model_cfg["unet_res_connect"] ) vae_cfg = OmegaConf.to_object(oc.model.vae) first_stage_model = AutoencoderKL( down_block_types=vae_cfg["down_block_types"], in_channels=vae_cfg["in_channels"], block_out_channels=vae_cfg["block_out_channels"], act_fn=vae_cfg["act_fn"], latent_channels=vae_cfg["latent_channels"], up_block_types=vae_cfg["up_block_types"], norm_num_groups=vae_cfg["norm_num_groups"], layers_per_block=vae_cfg["layers_per_block"], out_channels=vae_cfg["out_channels"], ) pretrained_ckpt_path = vae_cfg["pretrained_ckpt_path"] if pretrained_ckpt_path is not None: state_dict = torch.load(os.path.join(default_pretrained_vae_dir, vae_cfg["pretrained_ckpt_path"]), map_location=torch.device("cpu")) first_stage_model.load_state_dict(state_dict=state_dict) else: warnings.warn(f"Pretrained weights for `AutoencoderKL` not set. Run for sanity check only.") diffusion_cfg = OmegaConf.to_object(oc.model.diffusion) super(PreDiffSEVIRPLModule, self).__init__( torch_nn_module=latent_model, layout=oc.layout.layout, data_shape=diffusion_cfg["data_shape"], timesteps=diffusion_cfg["timesteps"], beta_schedule=diffusion_cfg["beta_schedule"], loss_type=self.oc.optim.loss_type, monitor=self.oc.optim.monitor, use_ema=diffusion_cfg["use_ema"], log_every_t=diffusion_cfg["log_every_t"], clip_denoised=diffusion_cfg["clip_denoised"], linear_start=diffusion_cfg["linear_start"], linear_end=diffusion_cfg["linear_end"], cosine_s=diffusion_cfg["cosine_s"], given_betas=diffusion_cfg["given_betas"], original_elbo_weight=diffusion_cfg["original_elbo_weight"], v_posterior=diffusion_cfg["v_posterior"], l_simple_weight=diffusion_cfg["l_simple_weight"], parameterization=diffusion_cfg["parameterization"], learn_logvar=diffusion_cfg["learn_logvar"], logvar_init=diffusion_cfg["logvar_init"], # latent diffusion latent_shape=diffusion_cfg["latent_shape"], first_stage_model=first_stage_model, cond_stage_model=diffusion_cfg["cond_stage_model"], num_timesteps_cond=diffusion_cfg["num_timesteps_cond"], cond_stage_trainable=diffusion_cfg["cond_stage_trainable"], cond_stage_forward=diffusion_cfg["cond_stage_forward"], scale_by_std=diffusion_cfg["scale_by_std"], scale_factor=diffusion_cfg["scale_factor"], ) # knowledge alignment knowledge_alignment_cfg = OmegaConf.to_object(oc.model.align) self.alignment_type = knowledge_alignment_cfg["alignment_type"] self.use_alignment = self.alignment_type is not None if self.use_alignment: alignment_ckpt_path = os.path.join(default_pretrained_alignment_dir, knowledge_alignment_cfg["model_ckpt_path"]) self.alignment_obj = SEVIRAvgIntensityAlignment( alignment_type=knowledge_alignment_cfg["alignment_type"], guide_scale=knowledge_alignment_cfg["guide_scale"], model_type=knowledge_alignment_cfg["model_type"], model_args=knowledge_alignment_cfg["model_args"], model_ckpt_path=alignment_ckpt_path ) disable_train(self.alignment_obj.model) self.alignment_model = self.alignment_obj.model alignment_fn = self.alignment_obj.get_mean_shift else: alignment_fn = None self.set_alignment(alignment_fn=alignment_fn) # lr_scheduler self.total_num_steps = total_num_steps # logging self.save_dir = save_dir self.logging_prefix = oc.logging.logging_prefix # visualization self.train_example_data_idx_list = list(oc.eval.train_example_data_idx_list) self.val_example_data_idx_list = list(oc.eval.val_example_data_idx_list) self.test_example_data_idx_list = list(oc.eval.test_example_data_idx_list) self.eval_example_only = oc.eval.eval_example_only if self.oc.eval.eval_unaligned: self.valid_mse = torchmetrics.MeanSquaredError() self.valid_mae = torchmetrics.MeanAbsoluteError() self.valid_score = SEVIRSkillScore( mode=self.oc.dataset.metrics_mode, seq_len=self.oc.layout.out_len, layout=self.layout, threshold_list=self.oc.dataset.threshold_list, metrics_list=self.oc.dataset.metrics_list, eps=1e-4 ) self.test_mse = torchmetrics.MeanSquaredError() self.test_mae = torchmetrics.MeanAbsoluteError() self.test_ssim = torchmetrics.image.StructuralSimilarityIndexMeasure() self.test_score = SEVIRSkillScore( mode=self.oc.dataset.metrics_mode, seq_len=self.oc.layout.out_len, layout=self.layout, threshold_list=self.oc.dataset.threshold_list, metrics_list=self.oc.dataset.metrics_list, eps=1e-4 ) self.test_fvd = FrechetVideoDistance( feature=self.oc.eval.fvd_features, layout=self.layout, reset_real_features=False, normalize=False, auto_t=True, ) if self.oc.eval.eval_aligned: self.valid_aligned_mse = torchmetrics.MeanSquaredError() self.valid_aligned_mae = torchmetrics.MeanAbsoluteError() self.valid_aligned_score = SEVIRSkillScore( mode=self.oc.dataset.metrics_mode, seq_len=self.oc.layout.out_len, layout=self.layout, threshold_list=self.oc.dataset.threshold_list, metrics_list=self.oc.dataset.metrics_list, eps=1e-4, ) self.test_aligned_mse = torchmetrics.MeanSquaredError() self.test_aligned_mae = torchmetrics.MeanAbsoluteError() self.test_aligned_ssim = torchmetrics.image.StructuralSimilarityIndexMeasure() self.test_aligned_score = SEVIRSkillScore( mode=self.oc.dataset.metrics_mode, seq_len=self.oc.layout.out_len, layout=self.layout, threshold_list=self.oc.dataset.threshold_list, metrics_list=self.oc.dataset.metrics_list, eps=1e-4, ) self.test_aligned_fvd = FrechetVideoDistance( feature=self.oc.eval.fvd_features, layout=self.layout, reset_real_features=False, normalize=False, auto_t=True, ) self.configure_save(cfg_file_path=oc_file) def configure_save(self, cfg_file_path=None): self.save_dir = os.path.join(default_exps_dir, self.save_dir) os.makedirs(self.save_dir, exist_ok=True) if cfg_file_path is not None: cfg_file_target_path = os.path.join(self.save_dir, "cfg.yaml") if (not os.path.exists(cfg_file_target_path)) or \ (not os.path.samefile(cfg_file_path, cfg_file_target_path)): copyfile(cfg_file_path, cfg_file_target_path) self.example_save_dir = os.path.join(self.save_dir, "examples") os.makedirs(self.example_save_dir, exist_ok=True) self.npy_save_dir = os.path.join(self.save_dir, "npy") os.makedirs(self.npy_save_dir, exist_ok=True) # region Get Default Config def get_base_config(self, oc_from_file=None): oc = OmegaConf.create() oc.layout = self.get_layout_config() oc.optim = self.get_optim_config() oc.logging = self.get_logging_config() oc.trainer = self.get_trainer_config() oc.eval = self.get_eval_config() oc.model = self.get_model_config() oc.dataset = self.get_dataset_config() if oc_from_file is not None: # oc = apply_omegaconf_overrides(oc, oc_from_file) oc = OmegaConf.merge(oc, oc_from_file) return oc @staticmethod def get_layout_config(): cfg = OmegaConf.create() cfg.in_len = 7 cfg.out_len = 6 cfg.in_step=1 cfg.out_step=1 cfg.in_out_diff=1 cfg.img_height = 128 cfg.img_width = 128 cfg.data_channels = 4 cfg.layout = "NTHWC" return cfg @staticmethod def get_model_config(): cfg = OmegaConf.create() layout_cfg = PreDiffSEVIRPLModule.get_layout_config() cfg.diffusion = OmegaConf.create() cfg.diffusion.data_shape = (layout_cfg.out_len, layout_cfg.img_height, layout_cfg.img_width, layout_cfg.data_channels) cfg.diffusion.timesteps = 1000 cfg.diffusion.beta_schedule = "linear" cfg.diffusion.use_ema = True cfg.diffusion.log_every_t = 100 # log every `log_every_t` timesteps. Must be smaller than `timesteps`. cfg.diffusion.clip_denoised = False cfg.diffusion.linear_start = 1e-4 cfg.diffusion.linear_end = 2e-2 cfg.diffusion.cosine_s = 8e-3 cfg.diffusion.given_betas = None cfg.diffusion.original_elbo_weight = 0. cfg.diffusion.v_posterior = 0. cfg.diffusion.l_simple_weight = 1. cfg.diffusion.parameterization = "eps" cfg.diffusion.learn_logvar = None cfg.diffusion.logvar_init = 0. # latent diffusion cfg.diffusion.latent_shape = [10, 16, 16, 4] cfg.diffusion.cond_stage_model = "__is_first_stage__" cfg.diffusion.num_timesteps_cond = None cfg.diffusion.cond_stage_trainable = False cfg.diffusion.cond_stage_forward = None cfg.diffusion.scale_by_std = False cfg.diffusion.scale_factor = 1.0 cfg.diffusion.latent_cond_shape = [10, 16, 16, 4] # knowledge alignment cfg.align = OmegaConf.create() cfg.align.alignment_type = None cfg.align.guide_scale = 1.0 cfg.align.model_type = "cuboid" cfg.align.model_ckpt_path = "tmp.pt" cfg.align.model_args = OmegaConf.create() # Earthformer cfg.align.model_args.input_shape = [6, 16, 16, 4] cfg.align.model_args.out_channels = 2 cfg.align.model_args.base_units = 16 cfg.align.model_args.block_units = None cfg.align.model_args.scale_alpha = 1.0 cfg.align.model_args.depth = [1, 1] cfg.align.model_args.downsample = 2 cfg.align.model_args.downsample_type = "patch_merge" cfg.align.model_args.block_attn_patterns = "axial" cfg.align.model_args.num_heads = 4 cfg.align.model_args.attn_drop = 0.0 cfg.align.model_args.proj_drop = 0.0 cfg.align.model_args.ffn_drop = 0.0 cfg.align.model_args.ffn_activation = "gelu" cfg.align.model_args.gated_ffn = False cfg.align.model_args.norm_layer = "layer_norm" cfg.align.model_args.use_inter_ffn = True cfg.align.model_args.hierarchical_pos_embed = False cfg.align.model_args.pos_embed_type = 't+h+w' cfg.align.model_args.padding_type = "zero" cfg.align.model_args.checkpoint_level = 0 cfg.align.model_args.use_relative_pos = True cfg.align.model_args.self_attn_use_final_proj = True # global vectors cfg.align.model_args.num_global_vectors = 0 cfg.align.model_args.use_global_vector_ffn = True cfg.align.model_args.use_global_self_attn = False cfg.align.model_args.separate_global_qkv = False cfg.align.model_args.global_dim_ratio = 1 # initialization cfg.align.model_args.attn_linear_init_mode = "0" cfg.align.model_args.ffn_linear_init_mode = "0" cfg.align.model_args.ffn2_linear_init_mode = "2" cfg.align.model_args.attn_proj_linear_init_mode = "2" cfg.align.model_args.conv_init_mode = "0" cfg.align.model_args.down_linear_init_mode = "0" cfg.align.model_args.global_proj_linear_init_mode = "2" cfg.align.model_args.norm_init_mode = "0" # timestep embedding for diffusion cfg.align.model_args.time_embed_channels_mult = 4 cfg.align.model_args.time_embed_use_scale_shift_norm = False cfg.align.model_args.time_embed_dropout = 0.0 # readout cfg.align.model_args.pool = "attention" cfg.align.model_args.readout_seq = True cfg.align.model_args.out_len = 6 cfg.latent_model = OmegaConf.create() cfg.latent_model.input_shape = [10, 16, 16, 4] cfg.latent_model.target_shape = [10, 16, 16, 4] cfg.latent_model.base_units = 4 # block_units = null cfg.latent_model.scale_alpha = 1.0 cfg.latent_model.num_heads = 4 cfg.latent_model.attn_drop = 0.1 cfg.latent_model.proj_drop = 0.1 cfg.latent_model.ffn_drop = 0.1 # inter-attn downsample/upsample cfg.latent_model.downsample = 2 cfg.latent_model.downsample_type = "patch_merge" cfg.latent_model.upsample_type = "upsample" cfg.latent_model.upsample_kernel_size = 3 # cuboid attention cfg.latent_model.depth = [1, 1] cfg.latent_model.self_pattern = "axial" # global vectors cfg.latent_model.num_global_vectors = 0 cfg.latent_model.use_dec_self_global = False cfg.latent_model.dec_self_update_global = True cfg.latent_model.use_dec_cross_global = False cfg.latent_model.use_global_vector_ffn = False cfg.latent_model.use_global_self_attn = True cfg.latent_model.separate_global_qkv = True cfg.latent_model.global_dim_ratio = 1 # mise cfg.latent_model.ffn_activation = "gelu" cfg.latent_model.gated_ffn = False cfg.latent_model.norm_layer = "layer_norm" cfg.latent_model.padding_type = "zeros" cfg.latent_model.pos_embed_type = "t+h+w" cfg.latent_model.checkpoint_level = 0 cfg.latent_model.use_relative_pos = True cfg.latent_model.self_attn_use_final_proj = True # initialization cfg.latent_model.attn_linear_init_mode = "0" cfg.latent_model.ffn_linear_init_mode = "0" cfg.latent_model.ffn2_linear_init_mode = "2" cfg.latent_model.attn_proj_linear_init_mode = "2" cfg.latent_model.conv_init_mode = "0" cfg.latent_model.down_up_linear_init_mode = "0" cfg.latent_model.global_proj_linear_init_mode = "2" cfg.latent_model.norm_init_mode = "0" # timestep embedding for diffusion cfg.latent_model.time_embed_channels_mult = 4 cfg.latent_model.time_embed_use_scale_shift_norm = False cfg.latent_model.time_embed_dropout = 0.0 cfg.latent_model.unet_res_connect = True cfg.vae = OmegaConf.create() cfg.vae.data_channels = layout_cfg.data_channels # from stable-diffusion-v1-5 cfg.vae.down_block_types = ['DownEncoderBlock2D', 'DownEncoderBlock2D', 'DownEncoderBlock2D', 'DownEncoderBlock2D'] cfg.vae.in_channels = cfg.vae.data_channels cfg.vae.block_out_channels = [128, 256, 512, 512] cfg.vae.act_fn = 'silu' cfg.vae.latent_channels = 4 cfg.vae.up_block_types = ['UpDecoderBlock2D', 'UpDecoderBlock2D', 'UpDecoderBlock2D', 'UpDecoderBlock2D'] cfg.vae.norm_num_groups = 32 cfg.vae.layers_per_block = 2 cfg.vae.out_channels = cfg.vae.data_channels return cfg @staticmethod def get_dataset_config(): cfg = OmegaConf.create() cfg.dataset_name = "sevir_lr" cfg.img_height = 128 cfg.img_width = 128 cfg.in_len = 7 cfg.out_len = 6 cfg.seq_len = 13 cfg.plot_stride = 1 cfg.interval_real_time = 10 cfg.sample_mode = "sequent" cfg.stride = cfg.out_len cfg.layout = "NTHWC" cfg.start_date = None cfg.train_val_split_date = (2019, 1, 1) cfg.train_test_split_date = (2019, 6, 1) cfg.end_date = None cfg.metrics_mode = "0" cfg.metrics_list = ('csi', 'pod', 'sucr', 'bias') cfg.threshold_list = (16, 74, 133, 160, 181, 219) cfg.aug_mode = "1" return cfg @staticmethod def get_optim_config(): cfg = OmegaConf.create() cfg.seed = None cfg.total_batch_size = 32 cfg.micro_batch_size = 8 cfg.float32_matmul_precision = "high" cfg.method = "adamw" cfg.lr = 1.0E-6 cfg.wd = 1.0E-2 cfg.betas = (0.9, 0.999) cfg.gradient_clip_val = 1.0 cfg.max_epochs = 50 cfg.loss_type = "l2" # scheduler cfg.warmup_percentage = 0.2 cfg.lr_scheduler_mode = "cosine" # Can be strings like 'linear', 'cosine', 'platue' cfg.min_lr_ratio = 1.0E-3 cfg.warmup_min_lr_ratio = 0.0 # early stopping cfg.monitor = "valid_loss_epoch" cfg.early_stop = False cfg.early_stop_mode = "min" cfg.early_stop_patience = 5 cfg.save_top_k = 1 return cfg @staticmethod def get_logging_config(): cfg = OmegaConf.create() cfg.logging_prefix = "PreDiff" cfg.monitor_lr = True cfg.monitor_device = False cfg.track_grad_norm = -1 cfg.use_wandb = False cfg.profiler = None cfg.save_npy = False return cfg @staticmethod def get_trainer_config(): cfg = OmegaConf.create() cfg.check_val_every_n_epoch = 1 cfg.log_step_ratio = 0.001 # Logging every 1% of the total training steps per epoch cfg.precision = 32 cfg.find_unused_parameters = True cfg.num_sanity_val_steps = 2 return cfg @staticmethod def get_eval_config(): cfg = OmegaConf.create() cfg.train_example_data_idx_list = [0, ] cfg.val_example_data_idx_list = [0, ] cfg.test_example_data_idx_list = [0, ] cfg.eval_example_only = False cfg.eval_aligned = True cfg.eval_unaligned = True cfg.num_samples_per_context = 1 cfg.font_size = 20 cfg.label_offset = (-0.5, 0.5) cfg.label_avg_int = False cfg.fvd_features = 400 return cfg # endregion # region Trainer and Optimizer Config def configure_optimizers(self): optim_cfg = self.oc.optim params = list(self.torch_nn_module.parameters()) if self.cond_stage_trainable: print(f"{self.__class__.__name__}: Also optimizing conditioner params!") params = params + list(self.cond_stage_model.parameters()) if self.learn_logvar: print('Diffusion model optimizing logvar') params.append(self.logvar) if optim_cfg.method == "adamw": optimizer = torch.optim.AdamW(params, lr=optim_cfg.lr, betas=optim_cfg.betas) else: raise NotImplementedError(f"opimization method {optim_cfg.method} not supported.") warmup_iter = int(np.round(self.oc.optim.warmup_percentage * self.total_num_steps)) if optim_cfg.lr_scheduler_mode == 'none': return {'optimizer': optimizer} else: if optim_cfg.lr_scheduler_mode == 'cosine': warmup_scheduler = LambdaLR(optimizer, lr_lambda=warmup_lambda(warmup_steps=warmup_iter, min_lr_ratio=optim_cfg.warmup_min_lr_ratio)) cosine_scheduler = CosineAnnealingLR(optimizer, T_max=(self.total_num_steps - warmup_iter), eta_min=optim_cfg.min_lr_ratio * optim_cfg.lr) lr_scheduler = SequentialLR(optimizer, schedulers=[warmup_scheduler, cosine_scheduler], milestones=[warmup_iter]) lr_scheduler_config = { 'scheduler': lr_scheduler, 'interval': 'step', 'frequency': 1, } else: raise NotImplementedError return {'optimizer': optimizer, 'lr_scheduler': lr_scheduler_config} def set_trainer_kwargs(self, **kwargs): r""" Default kwargs used when initializing pl.Trainer """ if self.oc.logging.profiler is None: profiler = None elif self.oc.logging.profiler == "pytorch": profiler = PyTorchProfiler(filename=f"{self.oc.logging.logging_prefix}_PyTorchProfiler.log") else: raise NotImplementedError checkpoint_callback = ModelCheckpoint( monitor=self.oc.optim.monitor, dirpath=os.path.join(self.save_dir, "checkpoints"), filename="{epoch:03d}_{val/loss:.4f}", auto_insert_metric_name=False, save_top_k=self.oc.optim.save_top_k, save_last=True, mode="min", ) callbacks = kwargs.pop("callbacks", []) assert isinstance(callbacks, list) for ele in callbacks: assert isinstance(ele, Callback) callbacks += [checkpoint_callback, ] if self.oc.logging.monitor_lr: callbacks += [LearningRateMonitor(logging_interval='step'), ] if self.oc.logging.monitor_device: callbacks += [DeviceStatsMonitor(), ] if self.oc.optim.early_stop: callbacks += [EarlyStopping(monitor=self.oc.optim.monitor, min_delta=0.0, patience=self.oc.optim.early_stop_patience, verbose=False, mode=self.oc.optim.early_stop_mode), ] logger = kwargs.pop("logger", []) tb_logger = pl_loggers.TensorBoardLogger(save_dir=self.save_dir) csv_logger = pl_loggers.CSVLogger(save_dir=self.save_dir) logger += [tb_logger, csv_logger] if self.oc.logging.use_wandb: wandb_logger = pl_loggers.WandbLogger( name = self.oc.logging.logging_name, id = self.oc.logging.run_id, project=self.oc.logging.logging_prefix, save_dir=self.save_dir ) logger += [wandb_logger, ] log_every_n_steps = max(1, int(self.oc.trainer.log_step_ratio * self.total_num_steps)) trainer_init_keys = inspect.signature(Trainer).parameters.keys() ret = dict( callbacks=callbacks, # log logger=logger, log_every_n_steps=log_every_n_steps, profiler=profiler, # save default_root_dir=self.save_dir, # ddp accelerator="gpu", strategy=DDPStrategy(find_unused_parameters=self.oc.trainer.find_unused_parameters), # strategy=ApexDDPStrategy(find_unused_parameters=False, delay_allreduce=True), # optimization max_epochs=self.oc.optim.max_epochs, check_val_every_n_epoch=self.oc.trainer.check_val_every_n_epoch, gradient_clip_val=self.oc.optim.gradient_clip_val, # NVIDIA amp precision=self.oc.trainer.precision, # misc num_sanity_val_steps=self.oc.trainer.num_sanity_val_steps, inference_mode=False, ) oc_trainer_kwargs = OmegaConf.to_object(self.oc.trainer) oc_trainer_kwargs = {key: val for key, val in oc_trainer_kwargs.items() if key in trainer_init_keys} ret.update(oc_trainer_kwargs) ret.update(kwargs) return ret # endregion # region Properties Extraction and Misc Calc @classmethod def get_total_num_steps( cls, num_samples: int, total_batch_size: int, epoch: int = None): r""" Parameters ---------- num_samples: int The number of samples of the datasets. `num_samples / micro_batch_size` is the number of steps per epoch. total_batch_size: int `total_batch_size == micro_batch_size * world_size * grad_accum` epoch: int """ if epoch is None: epoch = cls.get_optim_config().max_epochs return int(epoch * num_samples / total_batch_size) @staticmethod def get_sevir_datamodule(dataset_cfg, micro_batch_size: int = 1, num_workers: int = 4): dm = SEVIRLightningDataModule( seq_len=dataset_cfg["seq_len"], sample_mode=dataset_cfg["sample_mode"], stride=dataset_cfg["stride"], batch_size=micro_batch_size, layout=dataset_cfg["layout"], output_type=np.float32, preprocess=True, rescale_method="01", verbose=False, aug_mode=dataset_cfg["aug_mode"], ret_contiguous=False, # datamodule_only dataset_name=dataset_cfg["dataset_name"], start_date=dataset_cfg["start_date"], train_test_split_date=dataset_cfg["train_test_split_date"], end_date=dataset_cfg["end_date"], val_ratio=dataset_cfg["val_ratio"], num_workers=num_workers, ) return dm @property def in_slice(self): if not hasattr(self, "_in_slice"): in_slice, out_slice = step_layout_to_in_out_slice( layout=self.oc.layout.layout, in_len=self.oc.layout.in_len, in_step= self.oc.layout.in_step, out_len=self.oc.layout.out_len, out_step = self.oc.layout.out_step, in_out_diff= self.oc.layout.in_out_diff ) self._in_slice = in_slice self._out_slice = out_slice return self._in_slice @property def out_slice(self): if not hasattr(self, "_out_slice"): in_slice, out_slice = step_layout_to_in_out_slice( layout=self.oc.layout.layout, in_len=self.oc.layout.in_len, in_step= self.oc.layout.in_step, out_len=self.oc.layout.out_len, out_step = self.oc.layout.out_step, in_out_diff= self.oc.layout.in_out_diff ) self._in_slice = in_slice self._out_slice = out_slice return self._out_slice @torch.no_grad() def get_input(self, batch, **kwargs): r""" dataset dependent re-implement it for each specific dataset Parameters ---------- batch: Any raw data batch from specific dataloader Returns ------- out: Sequence[torch.Tensor, Dict[str, Any]] out[0] should be a torch.Tensor which is the target to generate out[1] should be a dict consists of several key-value pairs for conditioning """ return self._get_input_sevirlr(batch=batch, return_verbose=kwargs.get("return_verbose", False)) @torch.no_grad() def _get_input_sevirlr(self, batch, return_verbose=False): seq = batch in_seq = seq[self.in_slice] out_seq = seq[self.out_slice].contiguous() if return_verbose: return out_seq, {"y": in_seq}, in_seq else: return out_seq, {"y": in_seq} # endregion # region Operation Step def training_step(self, batch, batch_idx): loss, loss_dict = self(batch) self.log_dict(loss_dict, prog_bar=True, logger=True, on_step=True, on_epoch=True, sync_dist=False) micro_batch_size = batch.shape[self.batch_axis] data_idx = int(batch_idx * micro_batch_size) if self.current_epoch % self.oc.trainer.check_val_every_n_epoch == 0 \ and self.local_rank == 0: if data_idx in self.train_example_data_idx_list: target_seq, cond, context_seq = \ self.get_input(batch, return_verbose=True) aligned_pred_seq_list = [] aligned_pred_label_list = [] pred_seq_list = [] pred_label_list = [] for i in range(self.oc.eval.num_samples_per_context): # aligned sampling if self.use_alignment and self.oc.eval.eval_aligned: if self.alignment_type == "avg_x": alignment_kwargs = get_alignment_kwargs_avg_x(context_seq=context_seq, target_seq=target_seq) else: raise NotImplementedError pred_seq = self.sample( cond=cond, batch_size=micro_batch_size, return_intermediates=False, use_alignment=True, alignment_kwargs=alignment_kwargs, verbose=False, ).contiguous() aligned_pred_seq_list.append(pred_seq[0].detach().float().cpu().numpy()) aligned_pred_label_list.append(f"{self.oc.logging.logging_prefix}_aligned_pred_{i}") # no alignment if self.oc.eval.eval_unaligned: pred_seq = self.sample( cond=cond, batch_size=micro_batch_size, return_intermediates=False, verbose=False, ).contiguous() pred_seq_list.append(pred_seq[0].detach().float().cpu().numpy()) pred_label_list.append(f"{self.oc.logging.logging_prefix}_pred_{i}") pred_seq_list = aligned_pred_seq_list + pred_seq_list pred_label_list = aligned_pred_label_list + pred_label_list self.save_vis_step_end( data_idx=data_idx, context_seq=context_seq[0].detach().float().cpu().numpy(), target_seq=target_seq[0].detach().float().cpu().numpy(), pred_seq=pred_seq_list, pred_label=pred_label_list, mode="train", ) return loss def validation_step(self, batch, batch_idx): _, loss_dict_no_ema = self(batch) with self.ema_scope(): _, loss_dict_ema = self(batch) loss_dict_ema = {key + '_ema': loss_dict_ema[key] for key in loss_dict_ema} self.log_dict(loss_dict_no_ema, prog_bar=False, logger=True, on_step=False, on_epoch=True, sync_dist=True) self.log_dict(loss_dict_ema, prog_bar=False, logger=True, on_step=False, on_epoch=True, sync_dist=True) micro_batch_size = batch.shape[self.batch_axis] data_idx = int(batch_idx * micro_batch_size) if not self.eval_example_only or data_idx in self.val_example_data_idx_list: target_seq, cond, context_seq = \ self.get_input(batch, return_verbose=True) aligned_pred_seq_list = [] aligned_pred_label_list = [] pred_seq_list = [] pred_label_list = [] for i in range(self.oc.eval.num_samples_per_context): # aligned sampling if self.use_alignment and self.oc.eval.eval_aligned: if self.alignment_type == "avg_x": alignment_kwargs = get_alignment_kwargs_avg_x(context_seq=context_seq, target_seq=target_seq) else: raise NotImplementedError pred_seq = self.sample( cond=cond, batch_size=micro_batch_size, return_intermediates=False, use_alignment=True, alignment_kwargs=alignment_kwargs, verbose=False, ).contiguous() aligned_pred_seq_list.append(pred_seq[0].detach().float().cpu().numpy()) aligned_pred_label_list.append(f"{self.oc.logging.logging_prefix}_aligned_pred_{i}") if pred_seq.dtype is not torch.float: pred_seq = pred_seq.float() self.valid_aligned_mse(pred_seq, target_seq) self.valid_aligned_mae(pred_seq, target_seq) self.valid_aligned_score.update(pred_seq, target_seq) # no alignment if self.oc.eval.eval_unaligned: pred_seq = self.sample( cond=cond, batch_size=micro_batch_size, return_intermediates=False, verbose=False, ).contiguous() pred_seq_list.append(pred_seq[0].detach().float().cpu().numpy()) pred_label_list.append(f"{self.oc.logging.logging_prefix}_pred_{i}") if pred_seq.dtype is not torch.float: pred_seq = pred_seq.float() self.valid_mse(pred_seq, target_seq) self.valid_mae(pred_seq, target_seq) self.valid_score.update(pred_seq, target_seq) pred_seq_list = aligned_pred_seq_list + pred_seq_list pred_label_list = aligned_pred_label_list + pred_label_list self.save_vis_step_end( data_idx=data_idx, context_seq=context_seq[0].detach().float().cpu().numpy(), target_seq=target_seq[0].detach().float().cpu().numpy(), pred_seq=pred_seq_list, pred_label=pred_label_list, mode="val", suffix=f"_rank{self.local_rank}", ) def on_validation_epoch_end(self): if self.oc.eval.eval_unaligned: valid_mse = self.valid_mse.compute() valid_mae = self.valid_mae.compute() valid_score = self.valid_score.compute() valid_loss = -valid_score["avg"]["csi"] self.log('valid_loss_epoch', valid_loss, prog_bar=True, on_step=False, on_epoch=True, sync_dist=True) self.log('valid_mse_epoch', valid_mse, prog_bar=True, on_step=False, on_epoch=True, sync_dist=True) self.log('valid_mae_epoch', valid_mae, prog_bar=True, on_step=False, on_epoch=True, sync_dist=True) self.log_score_epoch_end(score_dict=valid_score, prefix="valid") self.valid_mse.reset() self.valid_mae.reset() self.valid_score.reset() if self.oc.eval.eval_aligned: valid_mse = self.valid_aligned_mse.compute() valid_mae = self.valid_aligned_mae.compute() valid_score = self.valid_aligned_score.compute() valid_loss = -valid_score["avg"]["csi"] self.log('valid_aligned_loss_epoch', valid_loss, prog_bar=True, on_step=False, on_epoch=True, sync_dist=True) self.log('valid_aligned_mse_epoch', valid_mse, prog_bar=True, on_step=False, on_epoch=True, sync_dist=True) self.log('valid_aligned_mae_epoch', valid_mae, prog_bar=True, on_step=False, on_epoch=True, sync_dist=True) self.log_score_epoch_end(score_dict=valid_score, prefix="valid_aligned") self.valid_aligned_mse.reset() self.valid_aligned_mae.reset() self.valid_aligned_score.reset() def test_step(self, batch, batch_idx): micro_batch_size = batch.shape[self.batch_axis] data_idx = int(batch_idx * micro_batch_size) if not self.eval_example_only or data_idx in self.val_example_data_idx_list: target_seq, cond, context_seq = \ self.get_input(batch, return_verbose=True) target_seq_bchw = rearrange(target_seq, "b t h w c -> (b t) c h w") aligned_pred_seq_list = [] aligned_pred_label_list = [] pred_seq_list = [] pred_label_list = [] for i in range(self.oc.eval.num_samples_per_context): # aligned sampling if self.use_alignment and self.oc.eval.eval_aligned: if self.alignment_type == "avg_x": alignment_kwargs = get_alignment_kwargs_avg_x(context_seq=context_seq, target_seq=target_seq) else: raise NotImplementedError pred_seq = self.sample( cond=cond, batch_size=micro_batch_size, return_intermediates=False, use_alignment=True, alignment_kwargs=alignment_kwargs, verbose=False, ).contiguous() if self.oc.logging.save_npy: npy_path = os.path.join(self.npy_save_dir, f"batch{batch_idx}_rank{self.local_rank}_sample{i}_aligned.npy") np.save(npy_path, pred_seq.detach().float().cpu().numpy()) aligned_pred_seq_list.append(pred_seq[0].detach().float().cpu().numpy()) aligned_pred_label_list.append(f"{self.oc.logging.logging_prefix}_aligned_pred_{i}") if pred_seq.dtype is not torch.float: pred_seq = pred_seq.float() self.test_aligned_mse(pred_seq, target_seq) self.test_aligned_mae(pred_seq, target_seq) self.test_aligned_score.update(pred_seq, target_seq) # self.test_aligned_fvd.update(pred_seq, real=False) pred_seq_bchw = rearrange(pred_seq, "b t h w c -> (b t) c h w") self.test_aligned_ssim(pred_seq_bchw, target_seq_bchw) # no alignment if self.oc.eval.eval_unaligned: pred_seq = self.sample( cond=cond, batch_size=micro_batch_size, return_intermediates=False, verbose=False, ).contiguous() if self.oc.logging.save_npy: npy_path = os.path.join(self.npy_save_dir, f"batch{batch_idx}_rank{self.local_rank}_sample{i}.npy") np.save(npy_path, pred_seq.detach().float().cpu().numpy()) pred_seq_list.append(pred_seq[0].detach().float().cpu().numpy()) pred_label_list.append(f"{self.oc.logging.logging_prefix}_pred_{i}") if pred_seq.dtype is not torch.float: pred_seq = pred_seq.float() self.test_mse(pred_seq, target_seq) self.test_mae(pred_seq, target_seq) self.test_score.update(pred_seq, target_seq) # self.test_fvd.update(pred_seq, real=False) pred_seq_bchw = rearrange(pred_seq, "b t h w c -> (b t) c h w") self.test_ssim(pred_seq_bchw, target_seq_bchw) # if self.use_alignment and self.oc.eval.eval_aligned: # self.test_aligned_fvd.update(target_seq, real=True) # if self.oc.eval.eval_unaligned: # self.test_fvd.update(target_seq, real=True) pred_seq_list = aligned_pred_seq_list + pred_seq_list pred_label_list = aligned_pred_label_list + pred_label_list self.save_vis_step_end( data_idx=data_idx, context_seq=context_seq[0].detach().float().cpu().numpy(), target_seq=target_seq[0].detach().float().cpu().numpy(), pred_seq=pred_seq_list, pred_label=pred_label_list, mode="test", suffix=f"_rank{self.local_rank}", ) def on_test_epoch_end(self): if self.oc.eval.eval_unaligned: test_mse = self.test_mse.compute() test_mae = self.test_mae.compute() test_ssim = self.test_ssim.compute() test_score = self.test_score.compute() # test_fvd = self.test_fvd.compute() self.log('test_mse_epoch', test_mse, prog_bar=True, on_step=False, on_epoch=True, sync_dist=True) self.log('test_mae_epoch', test_mae, prog_bar=True, on_step=False, on_epoch=True, sync_dist=True) self.log('test_ssim_epoch', test_ssim, prog_bar=True, on_step=False, on_epoch=True, sync_dist=True) self.log_score_epoch_end(score_dict=test_score, prefix="test") # self.log('test_fvd_epoch', test_fvd, prog_bar=True, on_step=False, on_epoch=True, sync_dist=True) self.test_mse.reset() self.test_mae.reset() self.test_ssim.reset() self.test_score.reset() # self.test_fvd.reset() if self.oc.eval.eval_aligned: test_mse = self.test_aligned_mse.compute() test_mae = self.test_aligned_mae.compute() test_ssim = self.test_aligned_ssim.compute() test_score = self.test_aligned_score.compute() # test_fvd = self.test_aligned_fvd.compute() self.log('test_aligned_mse_epoch', test_mse, prog_bar=True, on_step=False, on_epoch=True, sync_dist=True) self.log('test_aligned_mae_epoch', test_mae, prog_bar=True, on_step=False, on_epoch=True, sync_dist=True) self.log('test_aligned_ssim_epoch', test_ssim, prog_bar=True, on_step=False, on_epoch=True, sync_dist=True) self.log_score_epoch_end(score_dict=test_score, prefix="test_aligned") # self.log('test_aligned_fvd_epoch', test_fvd, prog_bar=True, on_step=False, on_epoch=True, sync_dist=True) self.test_aligned_mse.reset() self.test_aligned_mae.reset() self.test_aligned_ssim.reset() self.test_aligned_score.reset() # self.test_aligned_fvd.reset() # endregion def save_vis_step_end( self, data_idx: int, context_seq: np.ndarray, target_seq: np.ndarray, pred_seq: Union[np.ndarray, Sequence[np.ndarray]], pred_label: Union[str, Sequence[str]] = None, label_mode: str = "name", mode: str = "train", prefix: str = "", suffix: str = "", ): r""" Parameters ---------- data_idx context_seq, target_seq, pred_seq: np.ndarray layout should not include batch mode: str """ if mode == "train": example_data_idx_list = self.train_example_data_idx_list elif mode == "val": example_data_idx_list = self.val_example_data_idx_list elif mode == "test": example_data_idx_list = self.test_example_data_idx_list else: raise ValueError(f"Wrong mode {mode}! Must be in ['train', 'val', 'test'].") if label_mode == "name": # use the given label context_label = "context" target_label = "target" elif label_mode == "avg_int": context_label = f"context\navg_int={np.mean(context_seq):.4f}" target_label = f"target\navg_int={np.mean(target_seq):.4f}" if isinstance(pred_label, Sequence): pred_label = [f"{label}\navg_int={np.mean(seq):.4f}" for label, seq in zip(pred_label, pred_seq)] elif isinstance(pred_label, str): pred_label = f"{pred_label}\navg_int={np.mean(pred_seq):.4f}" else: raise TypeError(f"Wrong pred_label type {type(pred_label)}! must be in [str, Sequence[str]].") else: raise NotImplementedError if isinstance(pred_seq, Sequence): seq_list = [context_seq, target_seq] + list(pred_seq) label_list = [context_label, target_label] + pred_label else: seq_list = [context_seq, target_seq, pred_seq] label_list = [context_label, target_label, pred_label] if data_idx in example_data_idx_list: png_save_name = f"{prefix}{mode}_epoch_{self.current_epoch}_data_{data_idx}{suffix}.png" vis_sevir_seq( save_path=os.path.join(self.example_save_dir, png_save_name), seq=seq_list, label=label_list, interval_real_time=10, plot_stride=1, fs=self.oc.eval.fs, label_offset=self.oc.eval.label_offset, label_avg_int=self.oc.eval.label_avg_int, ) def log_score_epoch_end(self, score_dict: Dict, prefix: str = "valid"): for metrics in self.oc.dataset.metrics_list: for thresh in self.oc.dataset.threshold_list: score_mean = np.mean(score_dict[thresh][metrics]).item() self.log(f"{prefix}_{metrics}_{thresh}_epoch", score_mean, prog_bar=True, on_step=False, on_epoch=True, sync_dist=True) score_avg_mean = score_dict.get("avg", None) if score_avg_mean is not None: score_avg_mean = np.mean(score_avg_mean[metrics]).item() self.log(f"{prefix}_{metrics}_avg_epoch", score_avg_mean, prog_bar=True, on_step=False, on_epoch=True, sync_dist=True) def on_before_optimizer_step(self, optimizer): # Compute the 2-norm for each layer # If using mixed precision, the gradients are already unscaled here # reference: https://lightning.ai/docs/pytorch/2.0.9/debug/debugging_intermediate.html#look-out-for-exploding-gradients if self.oc.logging.track_grad_norm != -1: norms = grad_norm(self.torch_nn_module, norm_type=self.oc.logging.track_grad_norm) self.log_dict(norms)