NowcastNet¶

模型训练命令模型评估命令模型导出命令模型推理命令

暂无

# linux
wget -nc https://paddle-org.bj.bcebos.com/paddlescience/datasets/nowcastnet/mrms.tar
# windows
# curl https://paddle-org.bj.bcebos.com/paddlescience/datasets/nowcastnet/mrms.tar -o mrms.tar
mkdir ./datasets
tar -xvf mrms.tar -C ./datasets/
python nowcastnet.py mode=eval EVAL.pretrained_model_path=https://paddle-org.bj.bcebos.com/paddlescience/models/nowcastnet/nowcastnet_pretrained.pdparams

python nowcastnet.py mode=export

python nowcastnet.py mode=infer

1. 背景简介¶

近年来，深度学习方法已被应用于天气预报，尤其是雷达观测的降水预报。这些方法利用大量雷达复合观测数据来训练神经网络模型，以端到端的方式进行训练，无需明确参考降水过程的物理定律。这里复现了一个针对极端降水的非线性短临预报模型——NowcastNet，该模型将物理演变方案和条件学习法统一到一个神经网络框架中，实现了端到端的优化。

2. 模型原理¶

本章节仅对 NowcastNet 的模型原理进行简单地介绍，详细的理论推导请阅读 Skilful nowcasting of extreme precipitation with NowcastNet。

模型的总体结构如图所示：

模型使用预训练权重推理，接下来将介绍模型的推理过程。

3. 模型构建¶

在该案例中，用 PaddleScience 代码表示如下：

examples/nowcastnet/nowcastnet.py
if cfg.CASE_TYPE == "large":
    dataset_path = cfg.LARGE_DATASET_PATH
    model_cfg = cfg.MODEL.large
    output_dir = osp.join(cfg.output_dir, "large")
elif cfg.CASE_TYPE == "normal":
    dataset_path = cfg.NORMAL_DATASET_PATH
    model_cfg = cfg.MODEL.normal
    output_dir = osp.join(cfg.output_dir, "normal")
else:
    raise ValueError(
        f"cfg.CASE_TYPE should in ['normal', 'large'], but got '{cfg.mode}'"
    )
model = ppsci.arch.NowcastNet(**model_cfg)

examples/nowcastnet/conf/nowcastnet.yaml
CPU_WORKER: 1

# model settings
MODEL:
  normal:
    input_keys: ["input"]
    output_keys: ["output"]
    input_length: 9
    total_length: 29
    image_width: 512
    image_height: 512
    image_ch: 2
    ngf: 32
  large:
    input_keys: ["input"]
    output_keys: ["output"]
    input_length: 9
    total_length: 29
    image_width: 1024

其中，input_keys 和 output_keys 分别代表网络模型输入、输出变量的名称。

4. 模型评估可视化¶

完成上述设置之后，将上述实例化的对象按顺序传递给 ppsci.solver.Solver：

examples/nowcastnet/nowcastnet.py
solver = ppsci.solver.Solver(
    model,
    output_dir=output_dir,
    pretrained_model_path=cfg.EVAL.pretrained_model_path,
)

然后构建 VisualizerRadar 生成图片结果：

examples/nowcastnet/nowcastnet.py
visualizer = {
    "v_nowcastnet": ppsci.visualize.VisualizerRadar(
        {"input": frames_tensor},
        {
            "output": lambda out: out["output"],
        },
        prefix="v_nowcastnet",
        case_type=cfg.CASE_TYPE,
        total_length=model_cfg.total_length,
    )
}
solver.visualizer = visualizer
# visualize prediction
solver.visualize(batch_id)

5. 完整代码¶

examples/nowcastnet/nowcastnet.py
"""
Reference: https://codeocean.com/capsule/3935105/tree/v1
"""
from os import path as osp

import hydra
import paddle
from omegaconf import DictConfig

import ppsci
from ppsci.utils import logger


def train(cfg: DictConfig):
    print("Not supported.")


def evaluate(cfg: DictConfig):
    # set random seed for reproducibility
    ppsci.utils.misc.set_random_seed(cfg.seed)
    # initialize logger
    logger.init_logger("ppsci", osp.join(cfg.output_dir, "train.log"), "info")

    if cfg.CASE_TYPE == "large":
        dataset_path = cfg.LARGE_DATASET_PATH
        model_cfg = cfg.MODEL.large
        output_dir = osp.join(cfg.output_dir, "large")
    elif cfg.CASE_TYPE == "normal":
        dataset_path = cfg.NORMAL_DATASET_PATH
        model_cfg = cfg.MODEL.normal
        output_dir = osp.join(cfg.output_dir, "normal")
    else:
        raise ValueError(
            f"cfg.CASE_TYPE should in ['normal', 'large'], but got '{cfg.mode}'"
        )
    model = ppsci.arch.NowcastNet(**model_cfg)

    input_keys = ("radar_frames",)
    dataset_param = {
        "input_keys": input_keys,
        "label_keys": (),
        "image_width": model_cfg.image_width,
        "image_height": model_cfg.image_height,
        "total_length": model_cfg.total_length,
        "dataset_path": dataset_path,
        "data_type": paddle.get_default_dtype(),
    }
    test_data_loader = paddle.io.DataLoader(
        ppsci.data.dataset.RadarDataset(**dataset_param),
        batch_size=1,
        shuffle=False,
        num_workers=cfg.CPU_WORKER,
        drop_last=True,
    )

    # initialize solver
    solver = ppsci.solver.Solver(
        model,
        output_dir=output_dir,
        pretrained_model_path=cfg.EVAL.pretrained_model_path,
    )

    for batch_id, test_ims in enumerate(test_data_loader):
        test_ims = test_ims[0][input_keys[0]].numpy()
        frames_tensor = paddle.to_tensor(
            data=test_ims, dtype=paddle.get_default_dtype()
        )
        if batch_id <= cfg.NUM_SAVE_SAMPLES:
            visualizer = {
                "v_nowcastnet": ppsci.visualize.VisualizerRadar(
                    {"input": frames_tensor},
                    {
                        "output": lambda out: out["output"],
                    },
                    prefix="v_nowcastnet",
                    case_type=cfg.CASE_TYPE,
                    total_length=model_cfg.total_length,
                )
            }
            solver.visualizer = visualizer
            # visualize prediction
            solver.visualize(batch_id)


def export(cfg: DictConfig):
    from paddle.static import InputSpec

    # set models
    if cfg.CASE_TYPE == "large":
        model_cfg = cfg.MODEL.large
    elif cfg.CASE_TYPE == "normal":
        model_cfg = cfg.MODEL.normal
    else:
        raise ValueError(
            f"cfg.CASE_TYPE should in ['normal', 'large'], but got '{cfg.mode}'"
        )
    model = ppsci.arch.NowcastNet(**model_cfg)

    # load pretrained model
    solver = ppsci.solver.Solver(
        model=model, pretrained_model_path=cfg.INFER.pretrained_model_path
    )
    # export models
    input_spec = [
        {
            key: InputSpec(
                [None, 29, model_cfg.image_width, model_cfg.image_height, 2],
                "float32",
                name=key,
            )
            for key in model_cfg.input_keys
        },
    ]
    solver.export(input_spec, cfg.INFER.export_path)


def inference(cfg: DictConfig):
    import os.path as osp

    from deploy.python_infer import pinn_predictor

    # set model predictor
    predictor = pinn_predictor.PINNPredictor(cfg)

    if cfg.CASE_TYPE == "large":
        dataset_path = cfg.LARGE_DATASET_PATH
        model_cfg = cfg.MODEL.large
        output_dir = osp.join(cfg.output_dir, "large")
    elif cfg.CASE_TYPE == "normal":
        dataset_path = cfg.NORMAL_DATASET_PATH
        model_cfg = cfg.MODEL.normal
        output_dir = osp.join(cfg.output_dir, "normal")
    else:
        raise ValueError(
            f"cfg.CASE_TYPE should in ['normal', 'large'], but got '{cfg.mode}'"
        )

    input_keys = ("radar_frames",)
    dataset_param = {
        "input_keys": input_keys,
        "label_keys": (),
        "image_width": model_cfg.image_width,
        "image_height": model_cfg.image_height,
        "total_length": model_cfg.total_length,
        "dataset_path": dataset_path,
    }
    test_data_loader = paddle.io.DataLoader(
        ppsci.data.dataset.RadarDataset(**dataset_param),
        batch_size=cfg.INFER.batch_size,
        num_workers=cfg.CPU_WORKER,
        drop_last=True,
    )
    for batch_id, test_ims in enumerate(test_data_loader):
        if batch_id > cfg.NUM_SAVE_SAMPLES:
            break
        test_ims = {"input": test_ims[0][input_keys[0]].numpy()}
        output_dict = predictor.predict(test_ims, cfg.INFER.batch_size)
        # mapping data to model_cfg.output_keys
        output_dict = {
            store_key: output_dict[infer_key]
            for store_key, infer_key in zip(model_cfg.output_keys, output_dict.keys())
        }

        visualizer = ppsci.visualize.VisualizerRadar(
            test_ims,
            {
                "output": lambda out: out["output"],
            },
            prefix="v_nowcastnet",
            case_type=cfg.CASE_TYPE,
            total_length=model_cfg.total_length,
        )
        test_ims.update(output_dict)
        visualizer.save(osp.join(output_dir, f"epoch_{batch_id}"), test_ims)


@hydra.main(version_base=None, config_path="./conf", config_name="nowcastnet.yaml")
def main(cfg: DictConfig):
    if cfg.mode == "train":
        train(cfg)
    elif cfg.mode == "eval":
        evaluate(cfg)
    elif cfg.mode == "export":
        export(cfg)
    elif cfg.mode == "infer":
        inference(cfg)
    else:
        raise ValueError(
            f"cfg.mode should in ['train', 'eval', 'export', 'infer'], but got '{cfg.mode}'"
        )


if __name__ == "__main__":
    main()

6. 结果展示¶

下图展示了模型的预测结果和真值结果。

7. 参考资料¶

Skilful nowcasting of extreme precipitation with NowcastNet