import sys from pathlib import Path # 获取项目根目录（train.py上级的上级） root_path = Path(__file__).parent.parent sys.path.append(str(root_path)) import logging import os import time import numpy as np import torch import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel from onescience.datapipes.climate import CMEMSDatapipe from model.xihe import Xihe from onescience.utils.YParams import YParams from onescience.utils.fcn.darcy_loss import LpLoss def main(): logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s") logger = logging.getLogger() config_file_path = os.path.join(current_path, "conf/config.yaml") cfg = YParams(config_file_path, "model") 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() cfg_data = YParams(config_file_path, "datapipe") datapipe = CMEMSDatapipe( dataset_dir=cfg_data.dataset.data_dir, used_variables=cfg_data.dataset.channels, used_years=cfg_data.dataset.train_time, distributed=dist.is_initialized(), batch_size=cfg_data.dataloader.batch_size, num_workers=cfg_data.dataloader.num_workers, ) train_dataloader, train_sampler = datapipe.get_dataloader("train") datapipe = CMEMSDatapipe( dataset_dir=cfg_data.dataset.data_dir, used_variables=cfg_data.dataset.channels, used_years=cfg_data.dataset.val_time, distributed=dist.is_initialized(), batch_size=cfg_data.dataloader.batch_size, num_workers=cfg_data.dataloader.num_workers, ) val_dataloader, val_sampler = datapipe.get_dataloader("valid") model = Xihe(config=cfg).to(local_rank) optimizer = torch.optim.AdamW( model.parameters(), lr=cfg.lr, betas=tuple(cfg.betas), weight_decay=cfg.weight_decay, ) scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau( optimizer, factor=0.2, patience=5, mode="min" ) loss_obj = LpLoss() os.makedirs(cfg.checkpoint_dir, exist_ok=True) train_loss_file = f"{cfg.checkpoint_dir}/trloss.npy" valid_loss_file = f"{cfg.checkpoint_dir}/valoss.npy" best_valid_loss = 1.0e6 best_loss_epoch = 0 train_losses = np.empty((0,), dtype=np.float32) valid_losses = np.empty((0,), dtype=np.float32) if cfg.world_size == 1: total_params = sum(p.numel() for p in model.parameters()) print("\n") print("-" * 50) print(f"📂 now params is {total_params}, {total_params / 1e6:.2f}M, {total_params / 1e9:.2f}B") print("-" * 50, "\n") if os.path.exists(f"{cfg.checkpoint_dir}/model_bak.pth"): if world_rank == 0: print("\n") print("-" * 50) print("✅ 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) valid_losses = np.load(valid_loss_file) if cfg.world_size > 1: model = DistributedDataParallel(model, device_ids=[local_rank], output_device=local_rank) world_rank == 0 and logger.info("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.0 start_time = time.time() for j, data in enumerate(train_dataloader): invar = data[0].to(local_rank, dtype=torch.float32) outvar = data[1].to(local_rank, dtype=torch.float32) outvar_pred = model(invar) loss = loss_obj(outvar, outvar_pred) optimizer.zero_grad() loss.backward() optimizer.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}" ) train_loss /= len(train_dataloader) model.eval() valid_loss = 0.0 val_start_time = time.time() with torch.no_grad(): for j, data in enumerate(val_dataloader): invar = data[0].to(local_rank, dtype=torch.float32) outvar = data[1].to(local_rank, dtype=torch.float32) outvar_pred = model(invar) loss = loss_obj(outvar, outvar_pred) if cfg.world_size > 1: loss_tensor = loss.detach().to(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}-{j + 1}/{len(val_dataloader)} " f"[cost {int((time.time() - val_start_time) // 60):02}:{int((time.time() - val_start_time) % 60):02}] " f"[{(time.time() - val_start_time) / (j + 1): .02f}s/{cfg_data.dataloader.batch_size}batch] " f"loss:{valid_loss / (j + 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 scheduler.step(valid_loss) 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) valid_losses = np.append(valid_losses, valid_loss) np.save(train_loss_file, train_losses) np.save(valid_loss_file, valid_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") 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()