File size: 5,944 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 | import warnings
from collections import OrderedDict
from omegaconf import OmegaConf
import os
import argparse
import yaml
import torch
from lightning.pytorch import Trainer, seed_everything
from .prediff_lightning_module import PreDiffSEVIRPLModule
from utils.path import default_pretrained_earthformerunet_dir,pretrained_sevirlr_earthformer_unet_dir
from utils.pl_checkpoint import pl_load
from datamodule import WeatherForecastDataModuleOld
pytorch_state_dict_name = "sevirlr_earthformerunet.pt"
from dotenv import load_dotenv
_ = load_dotenv('./.env')
def get_parser():
parser = argparse.ArgumentParser()
parser.add_argument('--save', default='tmp_sevirlr_diffusion', type=str)
parser.add_argument('--nodes', default=1, type=int,
help="Number of nodes in DDP training.")
parser.add_argument('--gpus', nargs='+', type=int,
help="Number of GPUS per node in DDP training.")
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.')
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_earthformerunet_dir,pretrained_sevirlr_earthformer_unet_dir)), "Pretrained weights for Earthformer Unet 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(PreDiffSEVIRPLModule.get_dataset_config())
micro_batch_size = 1
total_batch_size = int(micro_batch_size * args.nodes * 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)
dm = PreDiffSEVIRPLModule.get_sevir_datamodule(
dataset_cfg=dataset_cfg,
micro_batch_size=micro_batch_size,
num_workers=4
)
dm.prepare_data()
dm.setup()
accumulate_grad_batches = total_batch_size // (micro_batch_size * args.nodes * len(args.gpus))
total_num_steps = PreDiffSEVIRPLModule.get_total_num_steps(
epoch=max_epochs,
num_samples=dm.num_train_samples,
total_batch_size=total_batch_size,
)
pl_module = PreDiffSEVIRPLModule(
total_num_steps=total_num_steps,
save_dir=args.save,
oc_file=args.cfg)
trainer_kwargs = pl_module.set_trainer_kwargs(
devices=args.gpus,
num_nodes=args.nodes,
accumulate_grad_batches=accumulate_grad_batches,
)
trainer = Trainer(**trainer_kwargs)
if args.pretrained:
# load Earthformer-UNet
earthformerunet_ckpt_path = os.path.join(
default_pretrained_earthformerunet_dir,
pretrained_sevirlr_earthformer_unet_dir
)
state_dict = torch.load(
earthformerunet_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)
pl_ckpt = pl_load(path_or_url=ckpt_path,
map_location=torch.device("cpu"))
# pl_state_dict = pl_ckpt["state_dict"] # pl 1.x
pl_state_dict = pl_ckpt
model_kay = "torch_nn_module."
model_state_dict = OrderedDict()
for key, val in pl_state_dict.items():
if key.startswith(model_kay):
model_state_dict[key.replace(model_kay, "")] = val
pl_module.torch_nn_module.load_state_dict(model_state_dict)
trainer.test(model=pl_module,
datamodule=dm, )
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 the latent diffusion model
pl_ckpt = pl_load(path_or_url=trainer.checkpoint_callback.best_model_path,
map_location=torch.device("cpu"))
# pl_state_dict = pl_ckpt["state_dict"] # pl 1.x
pl_state_dict = pl_ckpt
model_kay = "torch_nn_module."
state_dict = OrderedDict()
unexpected_dict = OrderedDict()
for key, val in pl_state_dict.items():
if key.startswith(model_kay):
state_dict[key.replace(model_kay, "")] = val
else:
unexpected_dict[key] = val
torch.save(state_dict, os.path.join(pl_module.save_dir, "checkpoints", pytorch_state_dict_name))
# test
trainer.test(ckpt_path="best",
datamodule=dm)
if __name__ == "__main__":
main()
|