adapt extformer-moe val and inference on device sdaa

.gitignore

@@ -132,3 +132,10 @@ FETCH_HEAD
 
 # auto generated version file by setuptools_scm
 ppsci/_version.py
+
+
+# Ignore entire output_laplace2d output directory
+/examples/laplace/output_laplace2d/
+/examples/extformer_moe/data
+/examples/extformer_moe/pretrained
+/examples/test

README.md

@@ -1,5 +1,7 @@
 # PaddleScience
 
+TecOrigin学习专用
+
 <!-- --8<-- [start:status] -->
 ![paddlescience_icon](https://paddle-org.bj.bcebos.com/paddlescience%2Fdocs%2Fpaddlescience_icon.png)
 > *Developed with [PaddlePaddle](https://www.paddlepaddle.org.cn/)*

examples/extformer_moe/conf/extformer_moe_enso_inference.yaml

@@ -0,0 +1,155 @@
+defaults:
+  - ppsci_default
+  - TRAIN: train_default
+  - TRAIN/ema: ema_default
+  - TRAIN/swa: swa_default
+  - EVAL: eval_default
+  - INFER: infer_default
+  - hydra/job/config/override_dirname/exclude_keys: exclude_keys_default
+  - _self_
+
+hydra:
+  run:
+    # dynamic output directory according to running time and override name
+    dir: outputs_extformer_moe_pretrain
+  job:
+    name: ${mode} # name of logfile
+    chdir: false # keep current working directory unchanged
+  callbacks:
+    init_callback:
+      _target_: ppsci.utils.callbacks.InitCallback
+  sweep:
+    # output directory for multirun
+    dir: ${hydra.run.dir}
+    subdir: ./
+
+# general settings
+mode: inference # running mode: inference only
+seed: 0
+output_dir: ${hydra:run.dir}
+log_freq: 20
+
+# set train and evaluate data path
+# FILE_PATH: /hpc2hdd/home/hni017/Workplace/data/weather_data/icar_enso_2021/enso_round1_train_20210201
+FILE_PATH: /localnvme/application/duff/workspace/PaddleScience/examples/extformer_moe/data/weather_data/icar_enso_2021/enso_round1_train_20210201/
+
+# dataset setting
+DATASET:
+  label_keys: ["sst_target","nino_target"]
+  in_len: 12
+  out_len: 14
+  nino_window_t: 3
+  in_stride: 1
+  out_stride: 1
+  train_samples_gap: 1
+  eval_samples_gap: 1
+  normalize_sst: true
+
+# model settings
+MODEL:
+  input_keys: ["sst_data"]
+  output_keys: ["sst_target","nino_target","aux_loss","rank_loss"]
+  input_shape: [12, 24, 48, 1]
+  target_shape: [14, 24, 48, 1]
+  base_units: 64
+  scale_alpha: 1.0
+
+  enc_depth: [1, 1]
+  dec_depth: [1, 1]
+  enc_use_inter_ffn: true
+  dec_use_inter_ffn: true
+  dec_hierarchical_pos_embed: false
+
+  downsample: 2
+  downsample_type: "patch_merge"
+  upsample_type: "upsample"
+
+  num_global_vectors: 0
+  use_dec_self_global: false
+  dec_self_update_global: true
+  use_dec_cross_global: false
+  use_global_vector_ffn: false
+  use_global_self_attn: false
+  separate_global_qkv: false
+  global_dim_ratio: 1
+
+  self_pattern: "axial"
+  cross_self_pattern: "axial"
+  cross_pattern: "cross_1x1"
+  dec_cross_last_n_frames: null
+
+  attn_drop: 0.1
+  proj_drop: 0.1
+  ffn_drop: 0.1
+  num_heads: 4
+
+  ffn_activation: "gelu"
+  gated_ffn: false
+  norm_layer: "layer_norm"
+  padding_type: "zeros"
+  pos_embed_type: "t+h+w"
+  use_relative_pos: true
+  self_attn_use_final_proj: true
+  dec_use_first_self_attn: false
+
+  z_init_method: "zeros"
+  initial_downsample_type: "conv"
+  initial_downsample_activation: "leaky_relu"
+  initial_downsample_scale: [1, 1, 2]
+  initial_downsample_conv_layers: 2
+  final_upsample_conv_layers: 1
+  checkpoint_level: 0
+
+  attn_linear_init_mode: "0"
+  ffn_linear_init_mode: "0"
+  conv_init_mode: "0"
+  down_up_linear_init_mode: "0"
+  norm_init_mode: "0"
+
+# moe settings
+MOE:
+  use_linear_moe: false
+  use_ffn_moe: true
+  use_attn_moe: false
+  num_experts: 10
+  out_planes: 4
+  importance_weight: 0.0
+  load_weight: 0.0
+  gate_style: "cuboid-latent" # linear, spatial-latent, cuboid-latent, spatial-latent-linear, cuboid-latent-linear
+  dispatch_style: "dense" # sparse, dense
+  aux_loss_style: "all" # all, cell
+
+# rnc settings
+RNC:
+  use_rnc: false
+  rank_imbalance_style: "batch+T+H+W"
+  feature_similarity_style: "l2"
+  rank_imbalance_temp: 2
+  label_difference_style: "l1"
+  rank_reg_coeff: 0.01
+  loss_cal_style: "computation-efficient" # computation-efficient, memory-efficient
+
+# training settings
+TRAIN:
+  epochs: 100
+  save_freq: 20
+  eval_during_train: true
+  eval_freq: 10
+  lr_scheduler:
+    epochs: ${TRAIN.epochs}
+    learning_rate: 0.0002
+    by_epoch: true
+  min_lr_ratio: 1.0e-3
+  wd: 1.0e-5
+  batch_size: 16
+  pretrained_model_path: null
+  checkpoint_path: null
+  update_freq: 1
+
+# evaluation settings
+EVAL:
+  # pretrained_model_path: ./checkpoint/enso/extformer_moe_enso.pdparams
+  pretrained_model_path: ./pretrained/extformer_moe_pretrained.pdparams
+  compute_metric_by_batch: false
+  eval_with_no_grad: true
+  batch_size: 1

examples/extformer_moe/conf/extformer_moe_enso_pretrain.yaml

@@ -30,7 +30,8 @@ output_dir: ${hydra:run.dir}
 log_freq: 20
 
 # set train and evaluate data path
-FILE_PATH: /hpc2hdd/home/hni017/Workplace/data/weather_data/icar_enso_2021/enso_round1_train_20210201
+# FILE_PATH: /hpc2hdd/home/hni017/Workplace/data/weather_data/icar_enso_2021/enso_round1_train_20210201
+FILE_PATH: /root/workspace/PaddleScience/examples/extformer_moe/data/weather_data/icar_enso_2021/enso_round1_train_20210201
 
 # dataset setting
 DATASET:

examples/extformer_moe/coverage.py

@@ -0,0 +1,29 @@
+import argparse
+import re
+
+def get_api_info(path):
+    pattern_fallback = '(?<=missing sdaa kernel: )\s*\w+'
+    pattern_api = '(?<=Finish AD API: )\s*\w+'
+    all_api_set = set()
+    fallback_api_set = set()
+    with open(path, 'r', encoding='utf-8') as f:
+        lines = f.readlines()
+    for line in lines:
+        if 'Finish AD API: ' in line:
+            api = re.search(pattern_api, line).group()
+            all_api_set.add(api)
+
+        if 'missing sdaa kernel: ' in line:
+            api = re.search(pattern_fallback, line).group()
+            fallback_api_set.add(api)
+    return sorted(list(all_api_set)), sorted(list(fallback_api_set))
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='算子覆盖率统计')
+    parser.add_argument('--path', type=str, help='日志文件路径', required=True)
+    args = parser.parse_args()
+    all_api_set, fallback_api_set = get_api_info(args.path)
+    print(f"all api: {all_api_set}, total: {len(all_api_set)}\n")
+    print(f"fallback op: {fallback_api_set}, total: {len(fallback_api_set)}\n")
+    print(f"coverage rate: {(1 - len(fallback_api_set) / len(all_api_set)) * 100:.2f}%")

examples/extformer_moe/extformer_moe_enso_inference.py

@@ -0,0 +1,77 @@
+# 需要加载环境变量export CUSTOM_DEVICE_BLACK_LIST=top_k_v2,top_k,mask_select
+import os
+import numpy as np
+import paddle
+from omegaconf import DictConfig
+from omegaconf import OmegaConf
+import  ppsci
+from ppsci.utils import save_load
+
+CMIP6_SST_MAX = 10.198975563049316
+CMIP6_SST_MIN = -16.549121856689453
+CMIP5_SST_MAX = 8.991744995117188
+CMIP5_SST_MIN = -9.33076286315918
+CMIP6_NINO_MAX = 4.138188362121582
+CMIP6_NINO_MIN = -3.5832221508026123
+CMIP5_NINO_MAX = 3.8253555297851562
+CMIP5_NINO_MIN = -2.691682815551758
+
+SST_MAX = max(CMIP6_SST_MAX, CMIP5_SST_MAX)
+SST_MIN = min(CMIP6_SST_MIN, CMIP5_SST_MIN)
+def scale_sst(sst):
+    return (sst - SST_MIN) / (SST_MAX - SST_MIN)
+
+
+
+
+def inference(cfg: DictConfig):
+    #  载入配置
+    normalize_sst = cfg.DATASET.normalize_sst
+    print('normalize_sst:', normalize_sst)
+    in_len = cfg.DATASET.in_len
+    print('in_len:', in_len)
+    input_keys = cfg.MODEL.input_keys
+
+    # 载入数据
+    test_00001_06_05 = np.load('./data/weather_data/icar_enso_2021/enso_final_test_data_B/test_00001_06_05.npy')
+
+    # 4为预测因子，并按照SST,T300,Ua,Va的顺序存放，所以只取sst
+    test_sst = paddle.to_tensor(test_00001_06_05[...,0], dtype='float32')
+    test_sst = test_sst[...,np.newaxis]
+    # 取19：67的区间是为了对应训练时的经度范围95E-330E
+    test_sst = test_sst[ :, :, 19:67, :]
+
+    test_sst_in_tar = np.concatenate([test_sst, test_sst, test_sst], axis=0)[:26]
+    print('test_sst:', test_sst_in_tar.shape)
+    if normalize_sst:
+        test_sst_in_tar = scale_sst(test_sst_in_tar)
+
+    # 构建模型所需的数据集
+    in_seq = paddle.unsqueeze(paddle.to_tensor(test_sst_in_tar[: in_len, ...], dtype='float32'), axis=0) # ( in_len, lat, lon, 1)
+    target_seq = paddle.unsqueeze(paddle.to_tensor(test_sst_in_tar[in_len :, ...], dtype='float32'), axis=0)  # ( in_len, lat, lon, 1)
+    input_item = {input_keys[0]: in_seq, "sst_target": target_seq}
+
+    # 载入模型
+    moe_config = OmegaConf.to_object(cfg.MOE)
+    rnc_config = OmegaConf.to_object(cfg.RNC)
+    model = ppsci.arch.ExtFormerMoECuboid(**cfg.MODEL, moe_config=moe_config, rnc_config=rnc_config)
+    save_load.load_pretrain(model, "./pretrained/extformer_moe_pretrained.pdparams")
+
+    # 预测
+    # model.eval()
+    pred = model(input_item)
+    print({k: (None if v is None else v.shape) for k, v in pred.items()})
+
+
+
+def main(cfg: DictConfig):
+   if cfg.mode == 'inference':
+        inference(cfg)
+   else:
+        raise ValueError("Invalid mode: for inference only, but got {}".format(cfg.mode))
+
+
+
+if __name__ == '__main__':
+    cfg = OmegaConf.load("./conf/extformer_moe_enso_inference.yaml")
+    main(cfg)

examples/extformer_moe/extformer_moe_enso_train.py

@@ -1,3 +1,4 @@
+# 需要加载环境变量export CUSTOM_DEVICE_BLACK_LIST=top_k_v2,top_k
 import enso_metric
 import hydra
 import paddle
@@ -7,6 +8,7 @@
 
 import ppsci
 
+paddle.set_device("sdaa")
 
 def get_parameter_names(model, forbidden_layer_types):
     result = []
@@ -43,7 +45,7 @@ def train(cfg: DictConfig):
             "shuffle": True,
         },
         "batch_size": cfg.TRAIN.batch_size,
-        "num_workers": 8,
+        "num_workers": 0,
     }
 
     # set constraint