File size: 6,230 Bytes
7667a87 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 | 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()))
# print(WeatherForecastDataModuleOld.__mro__)
# print(SEVIRLightningDataModule.__mro__)
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)
# save state_dict of VAE and discriminator
pl_ckpt = pl_load(path_or_url=trainer.checkpoint_callback.best_model_path,
map_location=torch.device("cpu"))
# state_dict = pl_ckpt["state_dict"] # pl 1.x
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))
# test
trainer.test(ckpt_path="best",
datamodule=dm)
if __name__ == "__main__":
main()
|