graph-cast-from-groud-zero/test/train_graphcast.py

#import hfai
#hfai.set_watchdog_time(21600)

import os
from pathlib import Path
import numpy as np
import torch
import torch.nn as nn
import torch.backends.cudnn as cudnn
import torch.distributed as dist
#import hfai.nccl.distributed as dist
#from haiscale.ddp import DistributedDataParallel
#from haiscale.pipeline import PipeDream, make_subgroups, partition
from torch.utils.data.distributed import DistributedSampler   # 干啥用的？
import timm.optim
from timm.scheduler import create_scheduler
import torch.multiprocessing as mp
from torch.utils.data import DataLoader
from model.graphcast_sequential import GraphCast



torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True

from data_factory.datasets import ERA5, EarthGraph
#from model.graphcast_sequential import get_graphcast_module
from utils.params import get_graphcast_args
from utils.tools import load_model, save_model
from utils.eval import graphcast_evaluate

SAVE_PATH = Path('/root/output/graphcast-torch/')
EarthGraph_PATH = Path('/root/code/OpenCastKit-removehfai/EarthGraph/')
SAVE_PATH.mkdir(parents=True, exist_ok=True)
EarthGraph_PATH.mkdir(parents=True, exist_ok=True)


def train_one_epoch(epoch, model, criterion, data_loader, graph, optimizer, lr_scheduler, min_loss, device):
    is_last_pipeline_stage = True   #(pp_group.rank() == pp_group.size() - 1)
    loss = torch.tensor(0., device="cuda")
    count = torch.tensor(0., device="cuda")
    model.train()

    input_x = [
        None,
        graph.mesh_data.x.half().cuda(non_blocking=True),
        graph.mesh_data.edge_index.cuda(non_blocking=True),
        graph.mesh_data.edge_attr.half().cuda(non_blocking=True),
        graph.grid2mesh_data.edge_index.cuda(non_blocking=True),
        graph.grid2mesh_data.edge_attr.half().cuda(non_blocking=True),
        graph.mesh2grid_data.edge_index.cuda(non_blocking=True),
        graph.mesh2grid_data.edge_attr.half().cuda(non_blocking=True)
    ]

    for step, batch in enumerate(data_loader):
        '''
        # Creates model and optimizer in default precision
        model = Net().cuda()
        optimizer = optim.SGD(model.parameters(), ...)

        for input, target in data:
            optimizer.zero_grad()

            # Enables autocasting for the forward pass (model + loss)
            with torch.autocast(device_type="cuda"):
                output = model(input)
                loss = loss_fn(output, target)

            # Exits the context manager before backward()
            loss.backward()
            optimizer.step()
        '''
        optimizer.zero_grad()
        
        x, y = [x.half().cuda(non_blocking=True) for x in batch]   # 在era5()中，已经给出 y，是两步预测。因为 banch =2 所以用两个数据进行预测
        input_x[0] = x   
        # 在这里，x 是多少都没有关系。因为 在 gx = self.grid_feat_embedding(gx) 特征嵌入的时候，怎样都可以。不管输入多少步，甚至输入的都不是应该的数据都行。
        # 不是的，x 必须和--grid-node-num 保持一致。因为         self.grid_feat_embedding = nn.Sequential(nn.Linear(gdim, gemb, bias=True) 决定了输入维度。

        with torch.cuda.amp.autocast():
            #step_loss, _ = model.forward_backward(*input_x, criterion=criterion, labels=(y,))
            out = model(input_x[0], input_x[1], input_x[2], input_x[3], input_x[4], input_x[5], input_x[6], input_x[7])  # *input_x  如果 “out = model(input_x)” 
            step_loss = criterion(out, y)

        loss.backward()
        optimizer.step()

        if is_last_pipeline_stage:
            loss += step_loss.sum().item()
            count += 1

        '''
        if dp_group.rank() == 0 and is_last_pipeline_stage and hfai.client.receive_suspend_command():
            save_model(model.module.module, epoch, step + 1, optimizer, lr_scheduler, min_loss, SAVE_PATH / 'latest.pt')
            hfai.go_suspend()
        '''
           
    # all-reduce in data paralel group
    if device == 0 and is_last_pipeline_stage:
        dist.all_reduce(loss)  # group=dp_group 不设置，使用默认
        dist.all_reduce(count)  # group=dp_group 不设置，使用默认
        loss = loss / count

    # broadcast from the last stage to other pipeline stages
    dist.all_reduce(loss)   # group=dp_group 不设置，使用默认

    return loss.item()


def train(local_rank, args):
    rank, world_size = dist.get_rank(), dist.get_world_size()

    # data parallel + pipeline parallel
    #dp_group, pp_group = make_subgroups(pp_size=args.pp_size)
    #dp_rank, dp_size = dp_group.rank(), dp_group.size()
    #pp_rank, pp_size = pp_group.rank(), pp_group.size()
    # is_last_pipeline_stage = (pp_group.rank() == pp_group.size() - 1)
    #dp_rank = # 默认就好，不用特别指定
    #dp_size =  # 训练线程数，感觉没啥用
    #print(f"RANK {rank}: data parallel {dp_rank}/{dp_size}", flush=True)


    # model & criterion & optimizer
    print("Load model & criterion & optimizer...")
    '''
    model = get_graphcast_module(args)????
    # model = hfai.nn.to_hfai(model)
    balance = [1, 1, 1, 1, 1, 1, 1, 1]????
    model = partition(model, pp_group.rank(), pp_group.size(), balance=balance)????

    model = DistributedDataParallel(model.cuda(), process_group=dp_group)????
    model = PipeDream(model, args.chunks, process_group=pp_group)????
    '''
    model = GraphCast(args).cuda()


    #param_groups = timm.optim.optim_factory(model, args.weight_decay)
    optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, betas=(0.9, 0.95))
    lr_scheduler, _ = create_scheduler(args, optimizer)
    criterion = nn.MSELoss()

    # generate graph
    if os.path.exists( EarthGraph_PATH / "EarthGraph"):
        graph = torch.load(EarthGraph_PATH / "EarthGraph")
    else:
        graph = EarthGraph()
        graph.generate_graph()
        torch.save(graph, EarthGraph_PATH / "EarthGraph")

    # load grid data
    print("Load data...")
    train_dataset = ERA5(split="train", check_data=True, modelname='graphcast')
    train_datasampler = DistributedSampler(train_dataset, shuffle=True)
    #train_dataloader = train_dataset.loader(args.batch_size, sampler=train_datasampler, num_workers=8, pin_memory=True, drop_last=True)
    train_dataloader = DataLoader(dataset=train_dataset, 
                              batch_size=args.batch_size,
                              sampler=train_datasampler, drop_last=True, num_workers=8, pin_memory = True)
    
    # 这里有问题。并没有区分测试和验证集。
    val_dataset = ERA5(split="val", check_data=True, modelname='graphcast')
    #val_datasampler = DistributedSampler(val_dataset, num_replicas=dp_size, rank=dp_rank, shuffle=True)
    val_datasampler = DistributedSampler(val_dataset, shuffle=True)
    #val_dataloader = val_dataset.loader(args.batch_size, sampler=val_datasampler, num_workers=8, pin_memory=True, drop_last=False)
    val_dataloader = DataLoader(dataset=val_dataset, 
                              batch_size=args.batch_size,
                              sampler=train_datasampler, drop_last=False, num_workers=8, pin_memory = True)
    
    # load
    start_epoch, start_step, min_loss = load_model(model, optimizer, lr_scheduler, SAVE_PATH / 'latest.pt')
    if local_rank == 0:
        print(f"Start training for {args.epochs} epochs")

    for epoch in range(start_epoch, args.epochs):

        train_loss = train_one_epoch(epoch, model, criterion, train_dataloader, graph, optimizer, lr_scheduler, min_loss, device="cuda:0")
        lr_scheduler.step(epoch)

        val_loss = graphcast_evaluate(val_dataloader, graph, model, criterion, dp_group, pp_group)

        if True:
            print(f"Epoch {epoch} | Train loss: {train_loss:.6f}, Val loss: {val_loss:.6f}")

        if True:
            save_model(model.module.module, epoch + 1, optimizer, lr_scheduler, min_loss, SAVE_PATH / 'latest.pt')
            if val_loss < min_loss:
                min_loss = val_loss
                save_model(model.module.module, path=SAVE_PATH / 'best.pt', only_model=True)

    # synchronize all processes
    model.module.reducer.stop()
    dist.barrier()


def main(local_rank, args):
    # fix the seed for reproducibility
    torch.manual_seed(2023)
    np.random.seed(2023)
    cudnn.benchmark = True

    # init dist
    ip = os.environ.get("MASTER_ADDR", "127.0.0.1")
    port = os.environ.get("MASTER_PORT", "22568")
    hosts = int(os.environ.get("WORLD_SIZE", "1"))  # number of nodes
    rank = int(os.environ.get("RANK", "0"))  # node id
    gpus = torch.cuda.device_count()  # gpus per node

    dist.init_process_group(backend="nccl", init_method=f"tcp://{ip}:{port}", world_size=hosts * gpus, rank=rank * gpus + local_rank)
    torch.cuda.set_device(local_rank)

    print("TRAINING STARTED...")
    train(local_rank, args)


if __name__ == '__main__':
    args = get_graphcast_args()
    ngpus = torch.cuda.device_count()
    #hfai.multiprocessing.spawn(main, args=(args,), nprocs=ngpus, bind_numa=True)
    main(0, args)