| import warnings |
| from collections import OrderedDict |
| from omegaconf import OmegaConf |
| import os |
| import argparse |
|
|
| import torch |
| import lightning.pytorch as pl |
| from lightning.pytorch import Trainer, seed_everything |
| from lightning.pytorch.utilities import grad_norm |
| import yaml |
|
|
| from .vae_lightning_module import VAESEVIRPLModule |
| from datamodule import WeatherForecastDataModuleOld,SEVIRLightningDataModule |
| from utils.pl_checkpoint import pl_load |
| from utils.path import ( |
| pretrained_sevirlr_vae_name, |
| default_pretrained_vae_dir |
| ) |
|
|
| from dotenv import load_dotenv |
| _ = load_dotenv('./.env') |
|
|
| pytorch_state_dict_name = "sevirlr_vae.pt" |
| pytorch_loss_state_dict_name = "sevirlr_vae_loss.pt" |
|
|
| def get_parser(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument('--save', default='tmp_sevirlr', type=str) |
| parser.add_argument('--gpus', nargs = '+', type=int) |
| parser.add_argument('--cfg', default=None, type=str) |
| parser.add_argument('--test', action='store_true') |
| parser.add_argument('--ckpt_name', default=None, type=str, |
| help='The model checkpoint trained on SEVIR-LR.') |
| parser.add_argument('--pretrained', action='store_true', |
| help='Load pretrained checkpoints for test.') |
| parser.add_argument('--sevir', action='store_true', |
| help='SEVIR Data') |
| parser.add_argument('--nhabe', action='store_true', |
| help='NhaBe Data') |
| return parser |
|
|
|
|
| def main(): |
| parser = get_parser() |
| args = parser.parse_args() |
| if args.pretrained: |
| args.cfg = os.path.abspath(os.path.join(os.path.dirname(__file__), "cfg.yaml")) |
| assert os.path.exists(os.path.join(default_pretrained_vae_dir,pretrained_sevirlr_vae_name)), "Pretrained weights for VAE does not exist" |
| |
| if args.cfg is not None: |
| oc_from_file = OmegaConf.load(open(args.cfg, "r")) |
| dataset_cfg = OmegaConf.to_object(oc_from_file.dataset) |
| total_batch_size = oc_from_file.optim.total_batch_size |
| micro_batch_size = oc_from_file.optim.micro_batch_size |
| max_epochs = oc_from_file.optim.max_epochs |
| seed = oc_from_file.optim.seed |
| float32_matmul_precision = oc_from_file.optim.float32_matmul_precision |
| else: |
| dataset_cfg = OmegaConf.to_object(VAESEVIRPLModule.get_dataset_config()) |
| micro_batch_size = 1 |
| total_batch_size = int(micro_batch_size * len(args.gpus)) |
| max_epochs = None |
| seed = 0 |
| float32_matmul_precision = "high" |
| torch.set_float32_matmul_precision(float32_matmul_precision) |
| seed_everything(seed, workers=True) |
| |
| if args.sevir: |
| dm = VAESEVIRPLModule.get_sevir_datamodule( |
| dataset_cfg=dataset_cfg, |
| micro_batch_size=micro_batch_size, |
| num_workers=8 |
| ) |
| dm.prepare_data() |
| dm.setup() |
| accumulate_grad_batches = total_batch_size // (micro_batch_size * len(args.gpus)) |
| total_num_steps = VAESEVIRPLModule.get_total_num_steps( |
| epoch=max_epochs, |
| num_samples=dm.num_train_samples, |
| total_batch_size=total_batch_size, |
| ) |
| elif args.nhabe: |
| with open('MambaUnet_servir.yaml','r') as f: |
| config = yaml.safe_load(f) |
| dm = WeatherForecastDataModuleOld( |
| **config['data'] |
| ) |
| dm.prepare_data() |
| dm.setup(stage=None) |
| accumulate_grad_batches = 1 |
| total_num_steps = int(oc_from_file.optim.max_epochs * len(dm.train_dataloader())) |
|
|
| |
| |
| |
| pl_module = VAESEVIRPLModule( |
| total_num_steps=total_num_steps, |
| accumulate_grad_batches=accumulate_grad_batches, |
| save_dir=args.save, |
| oc_file=args.cfg) |
| trainer_kwargs = pl_module.set_trainer_kwargs(devices=args.gpus) |
| trainer = Trainer(**trainer_kwargs) |
| if args.pretrained: |
| vae_ckpt_path = os.path.join(default_pretrained_vae_dir, |
| pretrained_sevirlr_vae_name) |
| state_dict = torch.load(vae_ckpt_path, |
| map_location=torch.device("cpu")) |
| pl_module.torch_nn_module.load_state_dict(state_dict=state_dict) |
| trainer.test(model=pl_module, |
| datamodule=dm) |
| elif args.test: |
| if args.ckpt_name is not None: |
| ckpt_path = os.path.join(pl_module.save_dir, "checkpoints", args.ckpt_name) |
| else: |
| ckpt_path = None |
| trainer.test(model=pl_module, |
| datamodule=dm, |
| ckpt_path=ckpt_path) |
| else: |
| if args.ckpt_name is not None: |
| ckpt_path = os.path.join(pl_module.save_dir, "checkpoints", args.ckpt_name) |
| if not os.path.exists(ckpt_path): |
| warnings.warn(f"ckpt {ckpt_path} not exists! Start training from epoch 0.") |
| ckpt_path = None |
| else: |
| ckpt_path = None |
| trainer.fit(model=pl_module, |
| datamodule=dm, |
| ckpt_path=ckpt_path) |
| |
| pl_ckpt = pl_load(path_or_url=trainer.checkpoint_callback.best_model_path, |
| map_location=torch.device("cpu")) |
| |
| state_dict = pl_ckpt |
| vae_key = "torch_nn_module." |
| vae_state_dict = OrderedDict() |
| loss_key = "loss." |
| loss_state_dict = OrderedDict() |
| unexpected_dict = OrderedDict() |
| for key, val in state_dict.items(): |
| if key.startswith(vae_key): |
| vae_state_dict[key[len(vae_key):]] = val |
| elif key.startswith(loss_key): |
| loss_state_dict[key[len(loss_key):]] = val |
| else: |
| unexpected_dict[key] = val |
| torch.save(vae_state_dict, os.path.join(pl_module.save_dir, "checkpoints", pytorch_state_dict_name)) |
| torch.save(loss_state_dict, os.path.join(pl_module.save_dir, "checkpoints", pytorch_loss_state_dict_name)) |
| |
| trainer.test(ckpt_path="best", |
| datamodule=dm) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
