NVIDIA Earth-2 Correction Diffusion NIM 的先决条件#

以下是使用 NVIDIA Earth-2 Correction Diffusion (CorrDiff) NIM 的先决条件。

支持矩阵#

模型#

模型名称

发布者

CorrDiff US GEFS - HRRR

NVIDIA

硬件支持#

优化配置#

GPU 内存和磁盘空间的值以 Gb 为单位；磁盘空间包括容器和模型；Profile 指的是模型优化的目标。

GPU

GPU 内存 (GB)

精度

GPU 数量

磁盘空间

H100

80

FP32

1

64

A100

40 & 80

FP32

1

64

RTX A6000

48

FP32

1

64

L40S

48

FP32

1

64

非优化配置#

GPU 内存和磁盘空间的值以 Gb 为单位；磁盘空间包括容器和模型。

GPU

GPU 内存 (GB)

精度

GPU 数量

磁盘空间

任何具有足够 GPU 内存和计算能力 ≥ 8.0 的 NVIDIA GPU

40

FP32

1

64

对于 GPU 内存较低的模型，可能需要在运行时进行一些额外的配置。即，调整 NIM 的目标样本批大小。有关更多详细信息，请参阅配置文档。

软件#

  • 在支持 Docker 的操作系统上，安装 Docker - 最低版本：23.0.1

  • 安装 NVIDIA 驱动程序 - 最低版本：545。但是，如果您在数据中心 GPU 上运行（例如，T4 或任何其他数据中心 GPU），则可以使用 NVIDIA 驱动程序版本 470.57（或更高版本 R470）、525.85（或更高版本 R525）、535.86（或更高版本 R535）或 545.23（或更高版本 R545）。

  • 安装 和 配置 NVIDIA Container Toolkit - 最低版本：1.13.5

  • 通过运行以下命令验证您的容器运行时是否支持 NVIDIA GPU

docker run --rm --runtime=nvidia --gpus all ubuntu nvidia-smi

示例输出

+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 560.35.03              Driver Version: 560.35.03      CUDA Version: 12.6     |
|-----------------------------------------+------------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id          Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |
|                                         |                        |               MIG M. |
|=========================================+========================+======================|
|   0  NVIDIA A100-PCIE-40GB          Off |   00000000:41:00.0 Off |                    0 |
| N/A   37C    P0             35W /  250W |   40423MiB /  40960MiB |      0%      Default |
|                                         |                        |             Disabled |
+-----------------------------------------+------------------------+----------------------+

有关枚举多 GPU 系统的更多信息，请参阅 NVIDIA Container Toolkit 的 GPU 枚举文档。

NGC 帐户#

您必须拥有经过身份验证的 NGC (NVIDIA GPU Cloud) 帐户，并有权访问相应的 NIM 和模型检查点。使用以下步骤登录 NGC，并设置 NGC_API_KEY 环境变量，以便您可以拉取镜像。

  1. 在 NGC 上创建一个帐户。

  2. 生成 API 密钥。以下步骤需要您的 NGC API 密钥。

  3. 通过运行以下代码，使用 NGC 验证本地 Docker。有关更多详细信息，请参阅 NGC 身份验证文档。

    docker login nvcr.io
Username: $oauthtoken
Password: <NGC API key>

  4. 在您的 shell 中设置 NGC API 密钥环境变量。您需要设置 NGC_API_KEY 变量。

    export NGC_API_KEY=<NGC API key>

客户端 Python 依赖项#

此 NIM 需要一个输入，其中包含所需时间戳的几个气象场，形式为 numpy 数组。可以使用 Earth2Studio 包轻松获取此输入数据，如快速入门指南所示。为此，设置一个隔离的 Python 环境（Python ≥ 3.10，< 3.12）并使用以下命令安装所需的依赖项

pip install --upgrade pip
pip install earth2studio

重要

此 NIM 需要 Earth2Studio 版本 >= 0.3.0

在提供的指南中，将使用其他几个软件包与 NIM 交互并后处理结果

pip install requests
pip install matplotlib

建议为客户端代码/输入数据创建一个空的工作目录

mkdir corrdiff
cd corrdiff

下一步#

安装所有先决条件后，要开始使用，请参阅快速入门指南。

在本页