NVIDIA Earth-2 FourCastNet NIM 快速入门指南#

使用本文档开始使用 NVIDIA Earth-2 FourCastNet NIM。

重要提示

在您可以使用本文档之前，您必须满足所有前提条件。

启动 NIM#

  1. 使用以下命令拉取 NIM 容器。

    注意

    容器大小约为 34GB（未压缩），下载时间取决于互联网连接速度。

    docker pull nvcr.io/nim/nvidia/fourcastnet:1.0.0

  2. 使用以下命令运行 NIM 容器。此命令启动 NIM 容器并公开端口 8000，供用户与 NIM 交互。它将模型拉取到本地文件系统上。

    注意

    模型大小约为 6GB，下载时间取决于互联网连接速度。

    export NGC_API_KEY=<NGC API Key>

docker run --rm --runtime=nvidia --gpus all --shm-size 4g \
    -p 8000:8000 \
    -e NGC_API_KEY \
    -t nvcr.io/nim/nvidia/fourcastnet:1.0.0

    NIM 运行后，验证您是否看到类似于以下的输出

    I0829 05:18:42.152983 108 grpc_server.cc:2466] "Started GRPCInferenceService at 0.0.0.0:8001"
I0829 05:18:42.153277 108 http_server.cc:4638] "Started HTTPService at 0.0.0.0:8090"
I0829 05:18:42.222628 108 http_server.cc:320] "Started Metrics Service at 0.0.0.0:8002"

检查 NIM 运行状况#

Bash#

  1. 打开一个新的终端，保持当前终端运行已启动的服务。

  2. 在新的终端中，等待运行状况检查端点返回 {"status":"ready"} 后再继续。这可能需要几分钟时间。使用以下命令查询运行状况检查。

    curl -X 'GET' \
    'http://127.0.0.1:8000/v1/health/ready' \
    -H 'accept: application/json'

Python#

import requests

r = requests.get("http://127.0.0.1:8000/v1/health/ready")
if r.status_code == 200:
   print("NIM is healthy!")
else:
   print("NIM is not ready!")

获取输入数据#

NIM 中使用的默认配置文件是球谐神经算子，它使用 73 个 ERA5 通道作为输入。此模型的具体输入/输出通道可以在其模型卡上找到。在本指南中，Earth-2 Studio 用于获取模型所需的输入数据。以下 Python 代码从 ARCO ERA5 获取 FourCastNet NIM 的初始状态，并将其保存到 NumPy 文件中

import numpy as np
from datetime import datetime
from earth2studio.data import ARCO
from earth2studio.models.px.sfno import VARIABLES

ds = ARCO()
da = ds(time=datetime(2023, 1, 1), variable=VARIABLES)
np.save("fcn_inputs.npy", da.to_numpy()[None].astype('float32'))

这将创建一个大小为 (1,1,73,721,1440) 的 NumPy 数组，这对应于模型卡中指定的所需输入字段。更广泛地说，这些维度映射到（批大小、提前期、变量、纬度和经度），如API 文档中所讨论的。

提示

有关此特定数据源的更多详细信息，请参阅 ARCO 数据源 API 文档。

推理请求#

使用 FourCastNet NIM 执行推理可以是标准的同步 post 请求。有关 NIM API 规范的完整文档，请访问 API 文档页面。

Bash#

curl -X POST \
        -F "input_array=@fcn_inputs.npy" \
        -F "input_time=2023-01-01T00:00:00Z" \
        -F "simulation_length=4" \
        -o output.tar \
        http://127.0.0.1:8000/v1/infer

Python#

import requests

url = "http://127.0.0.1:8000/v1/infer"
files = {
    "input_array": ("input_array", open("fcn_inputs.npy", "rb")),
}
data = {
    "input_time": "2023-01-01T00:00:00Z",
    "simulation_length": 4,
}
headers = {
    "accept": "application/x-tar",
}
print("Sending post request to NIM")

r = requests.post(url, headers=headers, data=data, files=files, timeout=180)

if r.status_code != 200:
    raise Exception(r.content)
else:
    # Dump response to file
    with open("output.tar", "wb") as tar:
        tar.write(r.content)

结果#

结果位于 tar 文件中，可以使用以下命令进行浏览

tar -tvf output.tar

-rw-r--r-- 0/0       303166208 1970-01-01 00:00 000_000.npy
-rw-r--r-- 0/0       303166208 1970-01-01 00:00 006_000.npy
-rw-r--r-- 0/0       303166208 1970-01-01 00:00 012_000.npy
-rw-r--r-- 0/0       303166208 1970-01-01 00:00 018_000.npy
-rw-r--r-- 0/0       303166208 1970-01-01 00:00 024_000.npy

提示

tar 存档中填充了 NumPy 数组，这些数组的命名约定为 {lead time}_{batch index}.npy。

警告

对于大多数推理调用，重要的是为对此 NIM 的请求指定更长的超时时间。

流式响应#

FourCastNet NIM 生成时间序列，并在生成每个时间步时将其流式传输回客户端。当创建在每个时间步上运行后续流程的管道时，这可能非常有用。以下代码片段可用于从 NIM 访问内存中生成的每个时间步。

import io
import tarfile
from pathlib import Path

import numpy as np
import requests
import tqdm

url = "http://127.0.0.1:8000/v1/infer"
files = {
    "input_array": ("input_array", open("fcn_inputs.npy", "rb")),
}
data = {
    "input_time": "2023-01-01T00:00:00Z",
    "simulation_length": 4,
}
headers = {
    "accept": "application/x-tar",
}

pbar = tqdm.tqdm(range(data["simulation_length"] + 1), desc="Forecast steps")
with (
    requests.post(
        url,
        headers=headers,
        files=files,
        data=data,
        timeout=300,
        stream=True,
    )
) as resp:
    resp.raise_for_status()
    with tarfile.open(fileobj=resp.raw, mode="r|") as tar_stream:
        # Loop over file members from tar stream
        for member in tar_stream:
            arr_file = io.BytesIO()
            arr_file.write(tar_stream.extractfile(member).read())
            arr_file.seek(0)
            data = np.load(arr_file)
            # Array names are on the form <LEAD_TIME>_<BATCH_IDX>.npy
            arr_lead_time, arr_batch_idx = (
                int(x) for x in Path(member.name).stem.split("_", maxsplit=1)
            )

            pbar.write(f"Received data for lead {arr_lead_time} batch {arr_batch_idx}")
            pbar.write(f"Output numpy {member.name} with shape {data.shape}")
            pbar.update(1)
pbar.close()

结果后处理#

最后一步是可视化 NIM 生成的结果。虽然您可以使用许多不同的工具来创建不同的可视化效果，但本指南演示了如何使用 Matplotlib。

import tarfile
import io

import numpy as np
import matplotlib.pyplot as plt

fig, ax = plt.subplots(2, 2)
ax = ax.flatten()

with tarfile.open("output.tar") as tar:
    for i, member in enumerate(tar.getmembers()):
        arr_file = io.BytesIO()
        arr_file.write(tar.extractfile(member).read())
        arr_file.seek(0)
        data = np.load(arr_file)

        ax[i].set_title(f"Lead time: {6*i}hrs")
        ax[i].imshow(data[0,0,0], vmin=-10, vmax=10, cmap="coolwarm")
        if i == 3:
            break

plt.savefig("output.png")

结果是 10 米高度处预测风速场的图。

fcn output

本页内容