# set initial states for convlstmNUM_CONVLSTM=cfg.num_convlstm(h0,c0)=(paddle.randn((1,128,16,16)),paddle.randn((1,128,16,16)))initial_state=[]for_inrange(NUM_CONVLSTM):initial_state.append((h0,c0))# grid parameterstime_steps=cfg.TIME_STEPSdx=cfg.DX[0]/cfg.DX[1]steps=cfg.TIME_BATCH_SIZE+1effective_step=list(range(0,steps))num_time_batch=int(time_steps/cfg.TIME_BATCH_SIZE)functions.dt=cfg.DTfunctions.dx=dxfunctions.time_steps=cfg.TIME_STEPSfunctions.num_time_batch=num_time_batchmodel=ppsci.arch.PhyCRNet(dt=cfg.DT,step=steps,effective_step=effective_step,**cfg.MODEL)
TIME_BATCH_SIZE:1000# model settingsMODEL:input_channels:2hidden_channels:[8,32,128,128]input_kernel_size:[4,4,4,3]input_stride:[2,2,2,1]input_padding:[1,1,1,1]
def_transform_out(_in,_out):returnfunctions.transform_out(_in,_out,model)model.register_input_transform(functions.transform_in)model.register_output_transform(_transform_out)# use Burgers_2d_solver_HighOrder.py to generate datadata=scio.loadmat(cfg.DATA_PATH)uv=data["uv"]# [t,c,h,w]functions.uv=uv# generate input data(input_dict_train,label_dict_train,input_dict_val,label_dict_val,)=functions.Dataset(paddle.to_tensor(initial_state),paddle.to_tensor(uv[0:1,],dtype=paddle.get_default_dtype(),),).get()
"""PhyCRNet for solving spatiotemporal PDEsReference: https://github.com/isds-neu/PhyCRNet/"""importfunctionsimporthydraimportpaddleimportscipy.ioassciofromomegaconfimportDictConfigimportppscideftrain(cfg:DictConfig):# set initial states for convlstmNUM_CONVLSTM=cfg.num_convlstm(h0,c0)=(paddle.randn((1,128,16,16)),paddle.randn((1,128,16,16)))initial_state=[]for_inrange(NUM_CONVLSTM):initial_state.append((h0,c0))# grid parameterstime_steps=cfg.TIME_STEPSdx=cfg.DX[0]/cfg.DX[1]steps=cfg.TIME_BATCH_SIZE+1effective_step=list(range(0,steps))num_time_batch=int(time_steps/cfg.TIME_BATCH_SIZE)functions.dt=cfg.DTfunctions.dx=dxfunctions.time_steps=cfg.TIME_STEPSfunctions.num_time_batch=num_time_batchmodel=ppsci.arch.PhyCRNet(dt=cfg.DT,step=steps,effective_step=effective_step,**cfg.MODEL)def_transform_out(_in,_out):returnfunctions.transform_out(_in,_out,model)model.register_input_transform(functions.transform_in)model.register_output_transform(_transform_out)# use Burgers_2d_solver_HighOrder.py to generate datadata=scio.loadmat(cfg.DATA_PATH)uv=data["uv"]# [t,c,h,w]functions.uv=uv# generate input data(input_dict_train,label_dict_train,input_dict_val,label_dict_val,)=functions.Dataset(paddle.to_tensor(initial_state),paddle.to_tensor(uv[0:1,],dtype=paddle.get_default_dtype(),),).get()sup_constraint_pde=ppsci.constraint.SupervisedConstraint({"dataset":{"name":"NamedArrayDataset","input":input_dict_train,"label":label_dict_train,},"batch_size":1,"num_workers":0,"auto_collation":True,},ppsci.loss.FunctionalLoss(functions.train_loss_func),{"loss":lambdaout:out["loss"],},name="sup_train",)constraint_pde={sup_constraint_pde.name:sup_constraint_pde}sup_validator_pde=ppsci.validate.SupervisedValidator({"dataset":{"name":"NamedArrayDataset","input":input_dict_val,"label":label_dict_val,},"batch_size":1,"num_workers":0,"auto_collation":True,},ppsci.loss.FunctionalLoss(functions.val_loss_func),{"loss":lambdaout:out["loss"],},metric={"metric":ppsci.metric.FunctionalMetric(functions.metric_expr)},name="sup_valid",)validator_pde={sup_validator_pde.name:sup_validator_pde}# initialize optimizer & schedulerscheduler=ppsci.optimizer.lr_scheduler.Step(**cfg.TRAIN.lr_scheduler)()optimizer=ppsci.optimizer.Adam(scheduler)(model)# initialize solversolver=ppsci.solver.Solver(model,constraint_pde,cfg.output_dir,optimizer,scheduler,cfg.TRAIN.epochs,cfg.TRAIN.iters_per_epoch,validator=validator_pde,eval_with_no_grad=cfg.TRAIN.eval_with_no_grad,pretrained_model_path=cfg.TRAIN.pretrained_model_path,)# train modelsolver.train()# evaluate after finished trainingmodel.register_output_transform(functions.tranform_output_val)solver.eval()defevaluate(cfg:DictConfig):# set initial states for convlstmNUM_CONVLSTM=cfg.num_convlstm(h0,c0)=(paddle.randn((1,128,16,16)),paddle.randn((1,128,16,16)))initial_state=[]for_inrange(NUM_CONVLSTM):initial_state.append((h0,c0))# grid parameterstime_steps=cfg.TIME_STEPSdx=cfg.DX[0]/cfg.DX[1]steps=cfg.EVAL.TIME_BATCH_SIZE+1effective_step=list(range(0,steps))num_time_batch=int(time_steps/cfg.EVAL.TIME_BATCH_SIZE)functions.dt=cfg.DTfunctions.dx=dxfunctions.num_time_batch=num_time_batchmodel=ppsci.arch.PhyCRNet(dt=cfg.DT,step=steps,effective_step=effective_step,**cfg.MODEL)def_transform_out(_in,_out):returnfunctions.transform_out(_in,_out,model)model.register_input_transform(functions.transform_in)model.register_output_transform(_transform_out)# use the generated datadata=scio.loadmat(cfg.DATA_PATH)uv=data["uv"]# [t,c,h,w]functions.uv=uv_,_,input_dict_val,_=functions.Dataset(paddle.to_tensor(initial_state),paddle.to_tensor(uv[0:1,...],dtype=paddle.get_default_dtype()),).get()ppsci.utils.load_pretrain(model,cfg.EVAL.pretrained_model_path)model.register_output_transform(None)functions.output_graph(model,input_dict_val,cfg.output_dir,cfg.case_name)@hydra.main(version_base=None,config_path="./conf",config_name="burgers_equations.yaml")defmain(cfg:DictConfig):ifcfg.mode=="train":train(cfg)elifcfg.mode=="eval":evaluate(cfg)else:raiseValueError(f"cfg.mode should in ['train', 'eval'], but got '{cfg.mode}'")if__name__=="__main__":main()