import sys from pathlib import Path # 获取项目根目录（train.py上级的上级） root_path = Path(__file__).parent.parent sys.path.append(str(root_path)) import torch import os import sys import numpy as np import torch.distributed as dist import logging import time from torch.nn.parallel import DistributedDataParallel from torch.optim.lr_scheduler import SequentialLR, LinearLR, CosineAnnealingLR, LambdaLR from onescience.datapipes.climate import ERA5Datapipe from onescience.utils.YParams import YParams from onescience.modules.utils.graphcast.data_utils import StaticData from onescience.modules.utils.graphcast.graph_utils import deg2rad from model.graph_cast_net import GraphCastNet from onescience.modules.utils.graphcast.loss import GraphCastLossFunction from apex import optimizers def main(): logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s") logger = logging.getLogger() ## Model config init config_file_path = os.path.join(current_path, "conf/config.yaml") cfg = YParams(config_file_path, "model") ## Distributed config init cfg.world_size = 1 if "WORLD_SIZE" in os.environ: cfg.world_size = int(os.environ["WORLD_SIZE"]) world_rank = 0 local_rank = 0 if cfg.world_size > 1: dist.init_process_group(backend="nccl", init_method="env://") local_rank = int(os.environ["LOCAL_RANK"]) world_rank = dist.get_rank() ## DataLoader init cfg_data = YParams(config_file_path, "datapipe") datapipe = ERA5Datapipe( dataset_dir=cfg_data.dataset.data_dir, used_variables=cfg_data.dataset.channels, used_years=cfg_data.dataset.train_time, distributed=dist.is_initialized(), num_workers=0 ) train_dataloader, train_sampler = datapipe.get_dataloader("train") datapipe = ERA5Datapipe( dataset_dir=cfg_data.dataset.data_dir, used_variables=cfg_data.dataset.channels, used_years=cfg_data.dataset.val_time, distributed=dist.is_initialized(), num_workers=0 ) val_dataloader, val_sampler = datapipe.get_dataloader("valid") input_dim_grid_nodes = (len(cfg_data.dataset.channels) + cfg.use_cos_zenith + 4 * cfg.use_time_of_year_index) * (cfg.num_history + 1) + cfg.num_channels_static model = GraphCastNet(mesh_level=cfg.mesh_level, multimesh=cfg.multimesh, input_res=tuple(cfg_data.dataset.img_size), input_dim_grid_nodes=input_dim_grid_nodes, input_dim_mesh_nodes=3, input_dim_edges=4, output_dim_grid_nodes=len(cfg_data.dataset.channels), processor_type=cfg.processor_type, khop_neighbors=cfg.khop_neighbors, num_attention_heads=cfg.num_attention_heads, processor_layers=cfg.processor_layers, hidden_dim=cfg.hidden_dim, norm_type=cfg.norm_type, do_concat_trick=cfg.concat_trick, recompute_activation=cfg.recompute_activation, ) model_dtype = torch.bfloat16 if cfg.full_bf16 else torch.float32 model.set_checkpoint_encoder(cfg.checkpoint_encoder) model.set_checkpoint_decoder(cfg.checkpoint_decoder) model = model.to(dtype=model_dtype).to(local_rank) if hasattr(model, "module"): latitudes = model.module.latitudes longitudes = model.module.longitudes lat_lon_grid = model.module.lat_lon_grid else: latitudes = model.latitudes longitudes = model.longitudes lat_lon_grid = model.lat_lon_grid static_dir = os.path.join(cfg_data.dataset.data_dir, "static") static_data = StaticData(static_dir, latitudes, longitudes).get().to(device=local_rank) channels_list = [i for i in range(len(cfg_data.dataset.channels))] area = torch.abs(torch.cos(deg2rad(lat_lon_grid[:, :, 0]))) area /= torch.mean(area) area = area.to(dtype=torch.bfloat16 if cfg.full_bf16 else torch.float32).to(device=local_rank) criterion = GraphCastLossFunction(area, channels_list, cfg_data.dataset.dataset_metadata_path, cfg_data.dataset.time_diff_std_path) optimizer = optimizers.FusedAdam(model.parameters(), lr=cfg.lr, betas=(0.9, 0.95), adam_w_mode=True, weight_decay=0.1) scheduler1 = LinearLR(optimizer, start_factor=1e-3, end_factor=1.0, total_iters=cfg.num_iters_step1, ) scheduler2 = CosineAnnealingLR(optimizer, T_max=cfg.num_iters_step2, eta_min=0.0) scheduler3 = LambdaLR(optimizer, lr_lambda=lambda epoch: (cfg.lr_step3 / cfg.lr)) scheduler = SequentialLR(optimizer, schedulers=[scheduler1, scheduler2, scheduler3], milestones=[cfg.num_iters_step1, cfg.num_iters_step1 + cfg.num_iters_step2]) ## Train process init os.makedirs(cfg.checkpoint_dir, exist_ok=True) train_loss_file = f"{cfg.checkpoint_dir}/trloss.npy" best_valid_loss = 1.0e6 best_loss_epoch = 0 train_losses = np.empty((0,), dtype=np.float32) ## Get model params count if cfg.world_size == 1: total_params = sum(p.numel() for p in model.parameters()) print("\n\n") print("-" * 50) print(f"📂 now params is {total_params}, {total_params / 1e6:.2f}M, {total_params / 1e9:.2f}B") print("-" * 50, "\n") ## Load model weight if there exist well-trained model if os.path.exists(f"{cfg.checkpoint_dir}/model_bak.pth"): if world_rank == 0: print("\n\n") print("-" * 50) print(f"✅ There has a model weight, load and continue training...") print(f'If you want to train a new model, ensure there is no *.pth file in {cfg.checkpoint_dir}') print("-" * 50, "\n") ckpt = torch.load(f"{cfg.checkpoint_dir}/model_bak.pth", map_location=f'cuda:{local_rank}', weights_only=False) model.load_state_dict(ckpt["model_state_dict"]) optimizer.load_state_dict(ckpt["optimizer_state_dict"]) scheduler.load_state_dict(ckpt["scheduler_state_dict"]) best_valid_loss = ckpt["best_valid_loss"] best_loss_epoch = ckpt["best_loss_epoch"] train_losses = np.load(train_loss_file) ## Distributed model if cfg.world_size > 1: model = DistributedDataParallel(model, device_ids=[local_rank], output_device=local_rank) world_rank == 0 and logger.info(f"start training ...") for epoch in range(cfg.max_epoch): if dist.is_initialized(): train_sampler.set_epoch(epoch) val_sampler.set_epoch(epoch) model.train() train_loss = 0 start_time = time.time() for j, data in enumerate(train_dataloader): invar = data[0].to(device=local_rank) outvar = data[1].to(device=local_rank) cos_zenith = data[2].to(device=local_rank) in_idx = data[3].item() cos_zenith = torch.squeeze(cos_zenith, dim=2) cos_zenith = torch.clamp(cos_zenith, min=0.0) - 1.0 / torch.pi day_of_year, time_of_day = divmod(in_idx * cfg.dt, 24) normalized_day_of_year = torch.tensor((day_of_year / 365) * (np.pi / 2), dtype=torch.float32, device=local_rank) normalized_time_of_day = torch.tensor((time_of_day / (24 - cfg.dt)) * (np.pi / 2), dtype=torch.float32, device=local_rank) sin_day_of_year = torch.sin(normalized_day_of_year).expand(1, 1, 721, 1440) cos_day_of_year = torch.cos(normalized_day_of_year).expand(1, 1, 721, 1440) sin_time_of_day = torch.sin(normalized_time_of_day).expand(1, 1, 721, 1440) cos_time_of_day = torch.cos(normalized_time_of_day).expand(1, 1, 721, 1440) invar = torch.concat((invar, cos_zenith, static_data, sin_day_of_year, cos_day_of_year, sin_time_of_day, cos_time_of_day), dim=1) invar, outvar = invar.to(dtype=model_dtype), outvar.to(dtype=model_dtype) outvar_pred = model(invar) loss = criterion(outvar_pred, outvar) optimizer.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.grad_clip_norm) torch.cuda.nvtx.range_pop() optimizer.step() scheduler.step() train_loss += loss.item() if world_rank == 0: logger.info(f'Train: Epoch {epoch}-{j+1}/{len(train_dataloader)} ' f'[cost {int((time.time()-start_time) // 60):02}:{int((time.time()-start_time) % 60):02}] ' f'[{(time.time()-start_time)/(j+1): .02f}s/{cfg_data.dataloader.batch_size}batch] ' f'loss:{train_loss / (j+1): .04f}') if (j + 1) % cfg.val_freq == 0: model.eval() valid_loss = 0.0 with torch.no_grad(): start_time = time.time() for k, data in enumerate(val_dataloader): invar = data[0].to(device=local_rank) outvar = data[1].to(device=local_rank) cos_zenith = data[2].to(device=local_rank) in_idx = data[3].item() cos_zenith = torch.squeeze(cos_zenith, dim=2) cos_zenith = torch.clamp(cos_zenith, min=0.0) - 1.0 / torch.pi # [b, 2, h, w] outvar = outvar.to(dtype=model_dtype) loss = 0.0 for t in range(outvar.shape[1]): day_of_year, time_of_day = divmod(in_idx + t * cfg.dt, 24 // cfg.dt) normalized_day_of_year = torch.tensor((day_of_year / 365) * (np.pi / 2), dtype=torch.float32, device=local_rank) normalized_time_of_day = torch.tensor((time_of_day / (24 - cfg.dt)) * (np.pi / 2), dtype=torch.float32, device=local_rank) sin_day_of_year = torch.sin(normalized_day_of_year).expand(1, 1, 721, 1440) cos_day_of_year = torch.cos(normalized_day_of_year).expand(1, 1, 721, 1440) sin_time_of_day = torch.sin(normalized_time_of_day).expand(1, 1, 721, 1440) cos_time_of_day = torch.cos(normalized_time_of_day).expand(1, 1, 721, 1440) invar = torch.concat((invar, cos_zenith, static_data, sin_day_of_year, cos_day_of_year, sin_time_of_day, cos_time_of_day), dim=1) invar = invar.to(dtype=model_dtype) outpred = model(invar) invar = outpred loss += criterion(outpred, outvar[:, t]) loss /= outvar.shape[1] if cfg.world_size > 1: loss_tensor = loss.detach().to(local_rank) # torch.tensor(loss, device=local_rank) dist.all_reduce(loss_tensor) loss = loss_tensor.item() / cfg.world_size valid_loss += loss else: valid_loss += loss.item() if world_rank == 0: logger.info(f'Valid: Epoch {epoch}-{k+1}/{len(val_dataloader)} ' f'[cost {int((time.time()-start_time) // 60):02}:{int((time.time()-start_time) % 60):02}] ' f'[{(time.time()-start_time)/(k+1): .02f}s/{cfg_data.dataloader.batch_size}batch] ' f'loss:{valid_loss / (k+1): .04f}') valid_loss /= len(val_dataloader) is_save_ckp = False if valid_loss < best_valid_loss: best_valid_loss = valid_loss best_loss_epoch = epoch world_rank == 0 and save_checkpoint(model, optimizer, scheduler, best_valid_loss, best_loss_epoch, cfg.checkpoint_dir) is_save_ckp = True train_loss /= (j+1) if world_rank == 0: logger.info(f"Epoch [{epoch + 1}/{cfg.max_epoch}], " f"Train Loss: {train_loss:.4f}, " f"Valid Loss: {valid_loss:.4f}, " f"Best loss at Epoch: {best_loss_epoch + 1}" + (", saving checkpoint" if is_save_ckp else "") ) train_losses = np.append(train_losses, train_loss) np.save(train_loss_file, train_losses) if epoch - best_loss_epoch > cfg.patience: print(f"Loss has not decrease in {cfg.patience} epochs, stopping training...") exit() def save_checkpoint(model, optimizer, scheduler, best_valid_loss, best_loss_epoch, model_path): model_to_save = model.module if hasattr(model, "module") else model state = {"model_state_dict": model_to_save.state_dict(), "optimizer_state_dict": optimizer.state_dict(), "scheduler_state_dict": scheduler.state_dict(), "best_valid_loss": best_valid_loss, "best_loss_epoch": best_loss_epoch, } torch.save(state, f"{model_path}/model.pth") ### the weight file saving may interrupted due to DCU queue limit, get a backup to ensure there at least has one model os.system(f"mv {model_path}/model.pth {model_path}/model_bak.pth") if __name__ == "__main__": current_path = os.getcwd() sys.path.append(current_path) main()