Skip to content

GitLab

Commit bdd87fae authored by zhangwenbo's avatar zhangwenbo
Browse files

Initial commit: FourCastNet source code only

parents
Hide whitespace changes
Inline Side-by-side
Showing with 1176 additions and 0 deletions
.gitignore 0 → 100755
View file @ bdd87fae
data.tar.gz
data/
*.tar.gz
*.zip
AUTHORS 0 → 100755
View file @ bdd87fae
The code was authored by the following people:
Jaideep Pathak - NVIDIA Corporation
Shashank Subramanian - NERSC, Lawrence Berkeley National Laboratory
Peter Harrington - NERSC, Lawrence Berkeley National Laboratory
Sanjeev Raja - NERSC, Lawrence Berkeley National Laboratory
Ashesh Chattopadhyay - Rice University
Morteza Mardani - NVIDIA Corporation
Thorsten Kurth - NVIDIA Corporation
David Hall - NVIDIA Corporation
Zongyi Li - California Institute of Technology, NVIDIA Corporation
Kamyar Azizzadenesheli - Purdue University
Pedram Hassanzadeh - Rice University
Karthik Kashinath - NVIDIA Corporation
Animashree Anandkumar - California Institute of Technology, NVIDIA Corporation
LICENSE 0 → 100755
View file @ bdd87fae
#BSD 3-Clause License
#
#Copyright (c) 2022, FourCastNet authors
#All rights reserved.
#
#Redistribution and use in source and binary forms, with or without
#modification, are permitted provided that the following conditions are met:
#
#1. Redistributions of source code must retain the above copyright notice, this
# list of conditions and the following disclaimer.
#
#2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
#3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
#THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
#AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
#IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
#DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
#FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
#DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
#SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
#CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
#OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
#OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
#The code was authored by the following people:
#
#Jaideep Pathak - NVIDIA Corporation
#Shashank Subramanian - NERSC, Lawrence Berkeley National Laboratory
#Peter Harrington - NERSC, Lawrence Berkeley National Laboratory
#Sanjeev Raja - NERSC, Lawrence Berkeley National Laboratory
#Ashesh Chattopadhyay - Rice University
#Morteza Mardani - NVIDIA Corporation
#Thorsten Kurth - NVIDIA Corporation
#David Hall - NVIDIA Corporation
#Zongyi Li - California Institute of Technology, NVIDIA Corporation
#Kamyar Azizzadenesheli - Purdue University
#Pedram Hassanzadeh - Rice University
#Karthik Kashinath - NVIDIA Corporation
#Animashree Anandkumar - California Institute of Technology, NVIDIA Corporation
README.md 0 → 100755
View file @ bdd87fae
# Fourcastnet_train
## 项目简介
---
## 环境部署
### 1. 拉取镜像
```bash
docker pull harbor.sourcefind.cn:5443/dcu/admin/base/pytorch:2.5.1-ubuntu22.04-dtk25.04.4-1230-py3.10-20260115
```
### 2. 创建容器
```bash
docker run -it \
--network=host \
--hostname=localhost \
--name=hunyuan \
-v /opt/hyhal:/opt/hyhal:ro \
-v $PWD:/workspace \
--ipc=host \
--device=/dev/kfd \
--device=/dev/mkfd \
--device=/dev/dri \
--shm-size=512G \
--privileged \
--group-add video \
--cap-add=SYS_PTRACE \
-u root \
--security-opt seccomp=unconfined \
harbor.sourcefind.cn:5443/dcu/admin/base/pytorch:2.5.1-ubuntu22.04-dtk25.04.4-1230-py3.10-20260115 \
/bin/bash
```
---
## 测试步骤
### 1. 拉取代码
```bash
git clone http://developer.sourcefind.cn/codes/bw-bestperf/hunyuanvideo-i2v.git
cd hunyuanvideo-i2v/
```
### 2. 安装依赖
```bash
pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install yunchang==0.6.0 xfuser==0.4.2
bash fix.sh # 适配xfuser
```
### 3. 下载模型
安装 ModelScope：
```bash
pip install modelscope
```
下载所需模型：
```bash
mkdir ckpts
modelscope download --model Tencent-Hunyuan/HunyuanVideo --local_dir ./ckpts
modelscope download --model AI-ModelScope/HunyuanVideo-I2V --local_dir ./ckpts
modelscope download --model AI-ModelScope/clip-vit-large-patch14 --local_dir ckpts/text_encoder_2
modelscope download --model AI-ModelScope/llava-llama-3-8b-v1_1-transformers --local_dir ckpts/text_encoder_i2v
```
执行修复脚本：
```bash
bash modified/fix.sh
```
设置环境变量，禁用 HIP 缓存 allocator 防止 OOM：
```bash
export PYTORCH_NO_HIP_MEMORY_CACHING=1
```
---
## 测试代码示例（四卡多GPU测试）
导出设备环境和禁用缓存：
```bash
export HIP_VISIBLE_DEVICES=4,5,6,7
export PYTORCH_NO_HIP_MEMORY_CACHING=1
```
运行文本到视频多GPU推理：
```bash
ALLOW_RESIZE_FOR_SP=1 torchrun --nproc_per_node=4 \
sample_image2video.py \
--model HYVideo-T/2 \
--prompt "An Asian man with short hair in black tactical uniform and white clothes waves a firework stick." \
--i2v-mode \
--i2v-image-path ./assets/demo/i2v/imgs/0.jpg \
--i2v-resolution 720p \
--i2v-stability \
--infer-steps 50 \
--video-length 129 \
--flow-reverse \
--flow-shift 17.0 \
--seed 42 \
--embedded-cfg-scale 6.0 \
--save-path ./results \
--ulysses-degree 4 \
--ring-degree 1 \
--num-videos 1 2>&1 | tee z_logs/bw_image2video_4ka.log
```
---
## 配置选项说明
| 参数 | 说明 | 默认值 / 示例 |
| -------------------- | ------------------------------ | ---------------------------------------- |
| `--model` | 指定使用的模型名称 | `HYVideo-T/2` |
| `--prompt` | 文本描述，用于生成视频 | `"An Asian man with short hair..."` |
| `--i2v-mode` | 启用文本到视频模式 | |
| `--i2v-image-path` | 输入图像路径 | `./assets/demo/i2v/imgs/0.jpg` |
| `--i2v-resolution` | 输出视频分辨率 | `720p` |
| `--i2v-stability` | 稳定性增强选项 | |
| `--infer-steps` | 推理步数，影响生成质量与速度 | `50` |
| `--video-length` | 生成视频的长度（帧数） | `129` |
| `--flow-reverse` | 是否反转光流 | |
| `--flow-shift` | 光流偏移值 | `17.0` |
| `--seed` | 随机种子，保证结果可复现 | `42` |
| `--embedded-cfg-scale` | Condition scaling比例 | `6.0` |
| `--save-path` | 生成结果保存路径 | `./results` |
| `--ulysses-degree` | 自定义参数，具体含义见代码说明 | `4` |
| `--ring-degree` | 自定义参数，具体含义见代码说明 | `1` |
| `--num-videos` | 生成视频数量 | `1` |
---
## 贡献指南
欢迎对 hunyuan-I2V 项目进行贡献！请遵循以下步骤：
1. Fork 本仓库，并新建分支进行功能开发或问题修复。
2. 提交规范的 commit 信息，描述清晰。
3. 提交 Pull Request，简述修改内容及目的。
4. 遵守项目代码规范和测试标准。
5. 参与代码评审，积极沟通改进方案。
---
## 许可证
本项目遵循 MIT 许可证，详见 [LICENSE](./LICENSE) 文件。
---
感谢您的关注与支持！如有问题，欢迎提交 Issue 或联系维护团队。
config/AFNO-50epoch.yaml 0 → 100755
View file @ bdd87fae
### base config ###
full_field: &FULL_FIELD
loss: 'l2'
lr: 1E-3
scheduler: 'ReduceLROnPlateau'
num_data_workers: 2
dt: 1 # how many timesteps ahead the model will predict
n_history: 0 #how many previous timesteps to consider
prediction_type: 'iterative'
prediction_length: 2 #applicable only if prediction_type == 'iterative'
n_initial_conditions: 5 #applicable only if prediction_type == 'iterative'
ics_type: "default"
save_raw_forecasts: !!bool True
save_channel: !!bool False
masked_acc: !!bool False
maskpath: None
perturb: !!bool False
add_grid: !!bool False
N_grid_channels: 0
gridtype: 'sinusoidal' #options 'sinusoidal' or 'linear'
roll: !!bool False
max_epochs: 50
batch_size: 64
#afno hyperparams
num_blocks: 8
nettype: 'afno'
patch_size: 8
width: 56
modes: 32
#options default, residual
target: 'default'
in_channels: [0,1]
out_channels: [0,1] #must be same as in_channels if prediction_type == 'iterative'
normalization: 'zscore' #options zscore (minmax not supported)
train_data_path: '/pscratch/sd/j/jpathak/wind/train'
valid_data_path: '/pscratch/sd/j/jpathak/wind/test'
inf_data_path: '/pscratch/sd/j/jpathak/wind/out_of_sample' # test set path for inference
exp_dir: '/pscratch/sd/j/jpathak/ERA5_expts_gtc/wind'
time_means_path: '/pscratch/sd/j/jpathak/wind/time_means.npy'
global_means_path: '/pscratch/sd/j/jpathak/wind/global_means.npy'
global_stds_path: '/pscratch/sd/j/jpathak/wind/global_stds.npy'
orography: !!bool False
orography_path: None
log_to_screen: !!bool True
log_to_wandb: !!bool False
save_checkpoint: !!bool True
enable_nhwc: !!bool False
optimizer_type: 'FusedAdam'
crop_size_x: None
crop_size_y: None
two_step_training: !!bool False
plot_animations: !!bool False
add_noise: !!bool False
noise_std: 0
afno_backbone: &backbone
<<: *FULL_FIELD
log_to_wandb: !!bool False
lr: 5E-4
batch_size: 8
max_epochs: 50
scheduler: 'CosineAnnealingLR'
in_channels: [0, 1 ,2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19]
out_channels: [0, 1 ,2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19]
orography: !!bool False
orography_path: None
exp_dir: './exp'
train_data_path: './data/train'
valid_data_path: './data/valid'
inf_data_path: './data/valid'
time_means_path: './data/time_means.npy'
global_means_path: './data/global_means.npy'
global_stds_path: './data/global_stds.npy'
afno_backbone_orography: &backbone_orography
<<: *backbone
orography: !!bool True
orography_path: '/pscratch/sd/s/shas1693/data/era5/static/orography.h5'
afno_backbone_finetune:
<<: *backbone
lr: 1E-4
batch_size: 64
log_to_wandb: !!bool False
max_epochs: 50
pretrained: !!bool True
two_step_training: !!bool True
pretrained_ckpt_path: '/pscratch/sd/s/shas1693/results/era5_wind/afno_backbone/0/training_checkpoints/best_ckpt.tar'
perturbations:
<<: *backbone
lr: 1E-4
batch_size: 64
max_epochs: 50
pretrained: !!bool True
two_step_training: !!bool True
pretrained_ckpt_path: '/pscratch/sd/j/jpathak/ERA5_expts_gtc/wind/afno_20ch_bs_64_lr5em4_blk_8_patch_8_cosine_sched/1/training_checkpoints/best_ckpt.tar'
prediction_length: 24
ics_type: "datetime"
n_perturbations: 100
save_channel: !bool True
save_idx: 4
save_raw_forecasts: !!bool False
date_strings: ["2018-01-01 00:00:00"]
inference_file_tag: " "
valid_data_path: "/pscratch/sd/j/jpathak/ "
perturb: !!bool True
n_level: 0.3
### PRECIP ###
precip: &precip
<<: *backbone
in_channels: [0, 1 ,2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19]
out_channels: [0]
nettype: 'afno'
nettype_wind: 'afno'
log_to_wandb: !!bool False
lr: 2.5E-4
batch_size: 64
max_epochs: 25
precip: '/pscratch/sd/p/pharring/ERA5/precip/total_precipitation'
time_means_path_tp: '/pscratch/sd/p/pharring/ERA5/precip/total_precipitation/time_means.npy'
model_wind_path: '/pscratch/sd/s/shas1693/results/era5_wind/afno_backbone_finetune/0/training_checkpoints/best_ckpt.tar'
precip_eps: !!float 1e-5
config/AFNO.yaml 0 → 100755
View file @ bdd87fae
### base config ###
full_field: &FULL_FIELD
loss: 'l2'
lr: 1E-3
scheduler: 'ReduceLROnPlateau'
num_data_workers: 4
dt: 1 # how many timesteps ahead the model will predict
n_history: 0 #how many previous timesteps to consider
prediction_type: 'iterative'
prediction_length: 41 #applicable only if prediction_type == 'iterative'
n_initial_conditions: 5 #applicable only if prediction_type == 'iterative'
ics_type: "default"
save_raw_forecasts: !!bool True
save_channel: !!bool False
masked_acc: !!bool False
maskpath: None
perturb: !!bool False
add_grid: !!bool False
N_grid_channels: 0
gridtype: 'sinusoidal' #options 'sinusoidal' or 'linear'
roll: !!bool False
max_epochs: 50
batch_size: 64
#afno hyperparams
num_blocks: 8
nettype: 'afno'
patch_size: 8
width: 56
modes: 32
#options default, residual
target: 'default'
in_channels: [0,1]
out_channels: [0,1] #must be same as in_channels if prediction_type == 'iterative'
normalization: 'zscore' #options zscore (minmax not supported)
train_data_path: '/pscratch/sd/j/jpathak/wind/train'
valid_data_path: '/pscratch/sd/j/jpathak/wind/test'
inf_data_path: '/pscratch/sd/j/jpathak/wind/out_of_sample' # test set path for inference
exp_dir: '/pscratch/sd/j/jpathak/ERA5_expts_gtc/wind'
time_means_path: '/pscratch/sd/j/jpathak/wind/time_means.npy'
global_means_path: '/pscratch/sd/j/jpathak/wind/global_means.npy'
global_stds_path: '/pscratch/sd/j/jpathak/wind/global_stds.npy'
orography: !!bool False
orography_path: None
log_to_screen: !!bool True
log_to_wandb: !!bool True
save_checkpoint: !!bool True
enable_nhwc: !!bool False
optimizer_type: 'FusedAdam'
crop_size_x: None
crop_size_y: None
two_step_training: !!bool False
plot_animations: !!bool False
add_noise: !!bool False
noise_std: 0
afno_backbone: &backbone
<<: *FULL_FIELD
log_to_wandb: !!bool True
lr: 5E-4
batch_size: 64
max_epochs: 150
scheduler: 'CosineAnnealingLR'
in_channels: [0, 1 ,2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19]
out_channels: [0, 1 ,2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19]
orography: !!bool False
orography_path: None
exp_dir: '/pscratch/sd/s/shas1693/results/era5_wind'
# train_data_path: '/pscratch/sd/s/shas1693/data/era5/train'
# valid_data_path: '/pscratch/sd/s/shas1693/data/era5/test'
# inf_data_path: '/pscratch/sd/s/shas1693/data/era5/out_of_sample'
# time_means_path: '/pscratch/sd/s/shas1693/data/era5/time_means.npy'
# global_means_path: '/pscratch/sd/s/shas1693/data/era5/global_means.npy'
# global_stds_path: '/pscratch/sd/s/shas1693/data/era5/global_stds.npy'
# ==== 数据路径（关键）====
train_data_path: '/workspace/FourCastNet/data/train'
valid_data_path: '/workspace/FourCastNet/data/train' # 先用同一份也可以跑
inf_data_path: '/workspace/FourCastNet/data/train'
# ==== 统计量（如果你现在还没有，先这样处理）====
time_means_path: '/workspace/FourCastNet/data/time_means.npy'
global_means_path: '/workspace/FourCastNet/data/global_means.npy'
global_stds_path: '/workspace/FourCastNet/data/global_stds.npy'
afno_backbone_orography: &backbone_orography
<<: *backbone
orography: !!bool True
orography_path: '/pscratch/sd/s/shas1693/data/era5/static/orography.h5'
afno_backbone_finetune:
<<: *backbone
lr: 1E-4
batch_size: 64
log_to_wandb: !!bool True
max_epochs: 50
pretrained: !!bool True
two_step_training: !!bool True
pretrained_ckpt_path: '/pscratch/sd/s/shas1693/results/era5_wind/afno_backbone/0/training_checkpoints/best_ckpt.tar'
perturbations:
<<: *backbone
lr: 1E-4
batch_size: 64
max_epochs: 50
pretrained: !!bool True
two_step_training: !!bool True
pretrained_ckpt_path: '/pscratch/sd/j/jpathak/ERA5_expts_gtc/wind/afno_20ch_bs_64_lr5em4_blk_8_patch_8_cosine_sched/1/training_checkpoints/best_ckpt.tar'
prediction_length: 24
ics_type: "datetime"
n_perturbations: 100
save_channel: !bool True
save_idx: 4
save_raw_forecasts: !!bool False
date_strings: ["2018-01-01 00:00:00"]
inference_file_tag: " "
valid_data_path: "/pscratch/sd/j/jpathak/ "
perturb: !!bool True
n_level: 0.3
### PRECIP ###
precip: &precip
<<: *backbone
in_channels: [0, 1 ,2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19]
out_channels: [0]
nettype: 'afno'
nettype_wind: 'afno'
log_to_wandb: !!bool True
lr: 2.5E-4
batch_size: 64
max_epochs: 25
precip: '/pscratch/sd/p/pharring/ERA5/precip/total_precipitation'
time_means_path_tp: '/pscratch/sd/p/pharring/ERA5/precip/total_precipitation/time_means.npy'
model_wind_path: '/pscratch/sd/s/shas1693/results/era5_wind/afno_backbone_finetune/0/training_checkpoints/best_ckpt.tar'
precip_eps: !!float 1e-5
copernicus/get_data_pl_short_length.py 0 → 100755
View file @ bdd87fae
import cdsapi
c = cdsapi.Client()
c.retrieve(
'reanalysis-era5-pressure-levels',
{
'product_type': 'reanalysis',
'format': 'netcdf',
'variable': [
'geopotential', 'relative_humidity', 'temperature',
'u_component_of_wind', 'v_component_of_wind',
],
'pressure_level': [
'50', '500', '850',
'1000',
],
'year': '2021',
'month': '10',
'day': [
'19', '20', '21',
'22', '23', '24',
'25', '26', '27',
'28', '29', '30',
'31',
],
'time': [
'00:00', '06:00', '12:00',
'18:00',
],
},
'/project/projectdirs/dasrepo/ERA5/oct_2021_19_31_pl.nc')
copernicus/get_data_sfc_short_length.py 0 → 100755
View file @ bdd87fae
import cdsapi
c = cdsapi.Client()
c.retrieve(
'reanalysis-era5-single-levels',
{
'product_type': 'reanalysis',
'format': 'netcdf',
'variable': [
'10m_u_component_of_wind', '10m_v_component_of_wind', '2m_temperature',
'mean_sea_level_pressure', 'surface_pressure', 'total_column_water_vapour',
],
'year': '2021',
'month': '10',
'day': [
'19', '20', '21',
'22', '23', '24',
'25', '26', '27',
'28', '29', '30',
'31',
],
'time': [
'00:00', '06:00', '12:00',
'18:00',
],
},
'/project/projectdirs/dasrepo/ERA5/oct_2021_19_31_sfc.nc')
# '/project/projectdirs/dasrepo/ERA5/oct_2021_19_31_pl.nc')
copernicus/get_data_u_v_6hr.py 0 → 100755
View file @ bdd87fae
import cdsapi
import numpy as np
import os
usr = 'j' # j,k,a
base_path = '/project/projectdirs/dasrepo/ERA5/wind_levels/6hr/' + usr
if not os.path.isdir(base_path):
os.makedirs(base_path)
year_dict = {'j': np.arange(1981, 1991), 'k': np.arange(1993,2006), 'a' : np.arange(2006, 2021)}
years = year_dict[usr]
#t1 = [str(jj).zfill(2) for jj in range(1,4)]
#t2 = [str(jj).zfill(2) for jj in range(4,7)]
#t3 = [str(jj).zfill(2) for jj in range(7,10)]
#t4 = [str(jj).zfill(2) for jj in range(10,13)]
#
#trimesters = [t1, t2, t3, t4]
#months = [str(jj).zfill(2) for jj in range(1,13)]
pressure_levels = [300]
c = cdsapi.Client()
for pressure_level in pressure_levels:
for year in years:
year_str = str(year)
pressure_str = str(pressure_level)
file_str = base_path + '/u_v_z_pressure_level_'+ pressure_str + '_' + year_str + '.nc'
print(year_str)
print(file_str)
c.retrieve(
'reanalysis-era5-pressure-levels',
{
'product_type': 'reanalysis',
'format': 'netcdf',
'pressure_level': pressure_str,
'variable': [
'u_component_of_wind', 'v_component_of_wind', 'geopotential',
],
'year': year_str,
'month': [
'01', '02', '03',
'04', '05', '06',
'07', '08', '09',
'10', '11', '12',
],
'day': [
'01', '02', '03',
'04', '05', '06',
'07', '08', '09',
'10', '11', '12',
'13', '14', '15',
'16', '17', '18',
'19', '20', '21',
'22', '23', '24',
'25', '26', '27',
'28', '29', '30',
'31',
],
'time': [
'00:00', '06:00', '12:00','18:00',
],
},
file_str)
data_process/get_stats.py 0 → 100755
View file @ bdd87fae
#BSD 3-Clause License
#
#Copyright (c) 2022, FourCastNet authors
#All rights reserved.
#
#Redistribution and use in source and binary forms, with or without
#modification, are permitted provided that the following conditions are met:
#
#1. Redistributions of source code must retain the above copyright notice, this
# list of conditions and the following disclaimer.
#
#2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
#3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
#THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
#AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
#IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
#DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
#FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
#DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
#SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
#CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
#OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
#OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
#The code was authored by the following people:
#
#Jaideep Pathak - NVIDIA Corporation
#Shashank Subramanian - NERSC, Lawrence Berkeley National Laboratory
#Peter Harrington - NERSC, Lawrence Berkeley National Laboratory
#Sanjeev Raja - NERSC, Lawrence Berkeley National Laboratory
#Ashesh Chattopadhyay - Rice University
#Morteza Mardani - NVIDIA Corporation
#Thorsten Kurth - NVIDIA Corporation
#David Hall - NVIDIA Corporation
#Zongyi Li - California Institute of Technology, NVIDIA Corporation
#Kamyar Azizzadenesheli - Purdue University
#Pedram Hassanzadeh - Rice University
#Karthik Kashinath - NVIDIA Corporation
#Animashree Anandkumar - California Institute of Technology, NVIDIA Corporation
import torch
import numpy as np
import h5py
years = [1979, 1989, 1999, 2004, 2010]
global_means = np.zeros((1,21,1,1))
global_stds = np.zeros((1,21,1,1))
time_means = np.zeros((1,21,721, 1440))
for ii, year in enumerate(years):
with h5py.File('/pscratch/sd/s/shas1693/data/era5/train/'+ str(year) + '.h5', 'r') as f:
rnd_idx = np.random.randint(0, 1460-500)
global_means += np.mean(f['fields'][rnd_idx:rnd_idx+500], keepdims=True, axis = (0,2,3))
global_stds += np.var(f['fields'][rnd_idx:rnd_idx+500], keepdims=True, axis = (0,2,3))
global_means = global_means/len(years)
global_stds = np.sqrt(global_stds/len(years))
time_means = time_means/len(years)
np.save('/pscratch/sd/s/shas1693/data/era5/global_means.npy', global_means)
np.save('/pscratch/sd/s/shas1693/data/era5/global_stds.npy', global_stds)
np.save('/pscratch/sd/s/shas1693/data/era5/time_means.npy', time_means)
print("means: ", global_means)
print("stds: ", global_stds)
data_process/normalize_orography.py 0 → 100755
View file @ bdd87fae
#BSD 3-Clause License
#
#Copyright (c) 2022, FourCastNet authors
#All rights reserved.
#
#Redistribution and use in source and binary forms, with or without
#modification, are permitted provided that the following conditions are met:
#
#1. Redistributions of source code must retain the above copyright notice, this
# list of conditions and the following disclaimer.
#
#2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
#3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
#THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
#AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
#IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
#DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
#FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
#DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
#SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
#CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
#OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
#OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
#The code was authored by the following people:
#
#Jaideep Pathak - NVIDIA Corporation
#Shashank Subramanian - NERSC, Lawrence Berkeley National Laboratory
#Peter Harrington - NERSC, Lawrence Berkeley National Laboratory
#Sanjeev Raja - NERSC, Lawrence Berkeley National Laboratory
#Ashesh Chattopadhyay - Rice University
#Morteza Mardani - NVIDIA Corporation
#Thorsten Kurth - NVIDIA Corporation
#David Hall - NVIDIA Corporation
#Zongyi Li - California Institute of Technology, NVIDIA Corporation
#Kamyar Azizzadenesheli - Purdue University
#Pedram Hassanzadeh - Rice University
#Karthik Kashinath - NVIDIA Corporation
#Animashree Anandkumar - California Institute of Technology, NVIDIA Corporation
import h5py
import numpy as np
with h5py.File('/pscratch/sd/s/shas1693/data/era5/static/orography.h5','a') as f:
orog = f['orog'][:]
omean = np.mean(orog)
print(omean)
ostd = np.std(orog)
print(ostd)
orog -= omean
orog /= ostd
f['orog'][...] = orog
f.flush()
data_process/parallel_copy.py 0 → 100755
View file @ bdd87fae
#BSD 3-Clause License
#
#Copyright (c) 2022, FourCastNet authors
#All rights reserved.
#
#Redistribution and use in source and binary forms, with or without
#modification, are permitted provided that the following conditions are met:
#
#1. Redistributions of source code must retain the above copyright notice, this
# list of conditions and the following disclaimer.
#
#2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
#3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
#THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
#AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
#IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
#DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
#FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
#DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
#SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
#CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
#OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
#OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
#The code was authored by the following people:
#
#Jaideep Pathak - NVIDIA Corporation
#Shashank Subramanian - NERSC, Lawrence Berkeley National Laboratory
#Peter Harrington - NERSC, Lawrence Berkeley National Laboratory
#Sanjeev Raja - NERSC, Lawrence Berkeley National Laboratory
#Ashesh Chattopadhyay - Rice University
#Morteza Mardani - NVIDIA Corporation
#Thorsten Kurth - NVIDIA Corporation
#David Hall - NVIDIA Corporation
#Zongyi Li - California Institute of Technology, NVIDIA Corporation
#Kamyar Azizzadenesheli - Purdue University
#Pedram Hassanzadeh - Rice University
#Karthik Kashinath - NVIDIA Corporation
#Animashree Anandkumar - California Institute of Technology, NVIDIA Corporation
import h5py
from mpi4py import MPI
import numpy as np
import time
from netCDF4 import Dataset as DS
import os
def writetofile(src, dest, channel_idx, varslist):
if os.path.isfile(src):
batch = 2**6
rank = MPI.COMM_WORLD.rank
Nproc = MPI.COMM_WORLD.size
Nimgtot = 1460#src_shape[0]
Nimg = Nimgtot//Nproc
print("Nimgtot",Nimgtot)
print("Nproc",Nproc)
print("Nimg",Nimg)
base = rank*Nimg
end = (rank+1)*Nimg if rank<Nproc - 1 else Nimgtot
idx = base
for variable_name in varslist:
fsrc = DS(src, 'r', format="NETCDF4").variables[variable_name]
fdest = h5py.File(dest, 'a', driver='mpio', comm=MPI.COMM_WORLD)
start = time.time()
while idx<end:
if end - idx < batch:
ims = fsrc[idx:end]
print(ims.shape)
fdest['fields'][idx:end, channel_idx, :, :] = ims
break
else:
ims = fsrc[idx:idx+batch]
fdest['fields'][idx:idx+batch, channel_idx, :, :] = ims
idx+=batch
ttot = time.time() - start
eta = (end - base)/((idx - base)/ttot)
hrs = eta//3600
mins = (eta - 3600*hrs)//60
secs = (eta - 3600*hrs - 60*mins)
ttot = time.time() - start
hrs = ttot//3600
mins = (ttot - 3600*hrs)//60
secs = (ttot - 3600*hrs - 60*mins)
channel_idx += 1
#year_dict = {'j': np.arange(1979, 1993), 'k': np.arange(1993, 2006), 'a' : np.arange(2006, 2021)}
dir_dict = {}
for year in np.arange(1979, 1993):
dir_dict[year] = 'j'
for year in np.arange(1993, 2006):
dir_dict[year] = 'k'
for year in np.arange(2006, 2021):
dir_dict[year] = 'a'
print(dir_dict)
years = np.arange(1979, 2018)
for year in years:
print(year)
src = '/global/cscratch1/sd/jpathak/ERA5_u10v10t2m_netcdf/10m_u_v_2m_t_gp_lsm_toa_' + str(year) + '.nc'
dest = '/global/cscratch1/sd/jpathak/ERA5/wind/vlevels/' + str(year) + '.h5'
writetofile(src, dest, 0, ['u'])
writetofile(src, dest, 1, ['v'])
src = '/project/projectdirs/dasrepo/jpathak/ERA5_more_data/z_u_v_1000_'+str(year)+'.nc'
writetofile(src, dest, 2, ['u'])
writetofile(src, dest, 3, ['v'])
writetofile(src, dest, 4, ['z'])
usr = dir_dict[year]
src ='/project/projectdirs/dasrepo/ERA5/wind_levels/6hr/' + usr + '/u_v_z_pressure_level_850_' +str(year) + '.nc'
writetofile(src, dest, 5, ['u'])
writetofile(src, dest, 6, ['v'])
writetofile(src, dest, 7, ['z'])
usr = dir_dict[year]
src ='/project/projectdirs/dasrepo/ERA5/wind_levels/6hr/' + usr + '/u_v_z_pressure_level_500_' +str(year) + '.nc'
writetofile(src, dest, 8, ['u'])
writetofile(src, dest, 9, ['v'])
writetofile(src, dest, 10, ['z'])
src = '/project/projectdirs/dasrepo/jpathak/ERA5_more_data/z50_'+str(year)+'.nc'
writetofile(src, dest, 11, ['z'])
data_process/parallel_copy_small_set.py 0 → 100755
View file @ bdd87fae
#BSD 3-Clause License
#
#Copyright (c) 2022, FourCastNet authors
#All rights reserved.
#
#Redistribution and use in source and binary forms, with or without
#modification, are permitted provided that the following conditions are met:
#
#1. Redistributions of source code must retain the above copyright notice, this
# list of conditions and the following disclaimer.
#
#2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
#3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
#THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
#AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
#IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
#DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
#FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
#DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
#SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
#CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
#OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
#OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
#The code was authored by the following people:
#
#Jaideep Pathak - NVIDIA Corporation
#Shashank Subramanian - NERSC, Lawrence Berkeley National Laboratory
#Peter Harrington - NERSC, Lawrence Berkeley National Laboratory
#Sanjeev Raja - NERSC, Lawrence Berkeley National Laboratory
#Ashesh Chattopadhyay - Rice University
#Morteza Mardani - NVIDIA Corporation
#Thorsten Kurth - NVIDIA Corporation
#David Hall - NVIDIA Corporation
#Zongyi Li - California Institute of Technology, NVIDIA Corporation
#Kamyar Azizzadenesheli - Purdue University
#Pedram Hassanzadeh - Rice University
#Karthik Kashinath - NVIDIA Corporation
#Animashree Anandkumar - California Institute of Technology, NVIDIA Corporation
# Instructions:
# Set Nimgtot correctly
import h5py
from mpi4py import MPI
import numpy as np
import time
from netCDF4 import Dataset as DS
import os
def writetofile(src, dest, channel_idx, varslist, src_idx=0, frmt='nc'):
if os.path.isfile(src):
batch = 2**4
rank = MPI.COMM_WORLD.rank
Nproc = MPI.COMM_WORLD.size
Nimgtot = 52#src_shape[0]
Nimg = Nimgtot//Nproc
base = rank*Nimg
end = (rank+1)*Nimg if rank<Nproc - 1 else Nimgtot
idx = base
for variable_name in varslist:
if frmt == 'nc':
fsrc = DS(src, 'r', format="NETCDF4").variables[variable_name]
elif frmt == 'h5':
fsrc = h5py.File(src, 'r')[varslist[0]]
print("fsrc shape", fsrc.shape)
fdest = h5py.File(dest, 'a', driver='mpio', comm=MPI.COMM_WORLD)
start = time.time()
while idx<end:
if end - idx < batch:
if len(fsrc.shape) == 4:
ims = fsrc[idx:end,src_idx]
else:
ims = fsrc[idx:end]
print(ims.shape)
fdest['fields'][idx:end, channel_idx, :, :] = ims
break
else:
if len(fsrc.shape) == 4:
ims = fsrc[idx:idx+batch,src_idx]
else:
ims = fsrc[idx:idx+batch]
#ims = fsrc[idx:idx+batch]
print("ims shape", ims.shape)
fdest['fields'][idx:idx+batch, channel_idx, :, :] = ims
idx+=batch
ttot = time.time() - start
eta = (end - base)/((idx - base)/ttot)
hrs = eta//3600
mins = (eta - 3600*hrs)//60
secs = (eta - 3600*hrs - 60*mins)
ttot = time.time() - start
hrs = ttot//3600
mins = (ttot - 3600*hrs)//60
secs = (ttot - 3600*hrs - 60*mins)
channel_idx += 1
filestr = 'oct_2021_19_31'
dest = '/global/cscratch1/sd/jpathak/21var/oct_2021_19_21.h5'
src = '/project/projectdirs/dasrepo/ERA5/oct_2021_19_31_sfc.nc'
#u10 v10 t2m
writetofile(src, dest, 0, ['u10'])
writetofile(src, dest, 1, ['v10'])
writetofile(src, dest, 2, ['t2m'])
#sp mslp
src = '/project/projectdirs/dasrepo/ERA5/oct_2021_19_31_sfc.nc'
writetofile(src, dest, 3, ['sp'])
writetofile(src, dest, 4, ['msl'])
#t850
src = '/project/projectdirs/dasrepo/ERA5/oct_2021_19_31_pl.nc'
writetofile(src, dest, 5, ['t'], 2)
#uvz1000
src = '/project/projectdirs/dasrepo/ERA5/oct_2021_19_31_pl.nc'
writetofile(src, dest, 6, ['u'], 3)
writetofile(src, dest, 7, ['v'], 3)
writetofile(src, dest, 8, ['z'], 3)
#uvz850
src = '/project/projectdirs/dasrepo/ERA5/oct_2021_19_31_pl.nc'
writetofile(src, dest, 9, ['u'], 2)
writetofile(src, dest, 10, ['v'], 2)
writetofile(src, dest, 11, ['z'], 2)
#uvz 500
src = '/project/projectdirs/dasrepo/ERA5/oct_2021_19_31_pl.nc'
writetofile(src, dest, 12, ['u'], 1)
writetofile(src, dest, 13, ['v'], 1)
writetofile(src, dest, 14, ['z'], 1)
#t500
src = '/project/projectdirs/dasrepo/ERA5/oct_2021_19_31_pl.nc'
writetofile(src, dest, 15, ['t'], 1)
#z50
src = '/project/projectdirs/dasrepo/ERA5/oct_2021_19_31_pl.nc'
writetofile(src, dest, 16, ['z'], 0)
#r500
src = '/project/projectdirs/dasrepo/ERA5/oct_2021_19_31_pl.nc'
writetofile(src, dest, 17, ['r'], 1)
#r850
src = '/project/projectdirs/dasrepo/ERA5/oct_2021_19_31_pl.nc'
writetofile(src, dest, 18, ['r'], 2)
#tcwv
src = '/project/projectdirs/dasrepo/ERA5/oct_2021_19_31_sfc.nc'
writetofile(src, dest, 19, ['tcwv'])
#sst
#src = '/project/projectdirs/dasrepo/ERA5/oct_2021_19_31_sfc.nc'
#writetofile(src, dest, 20, ['sst'])
docker/Dockerfile 0 → 100755
View file @ bdd87fae
FROM nvcr.io/nvidia/pytorch:21.11-py3
# update repo info
RUN apt update -y
# install mpi4py
RUN pip install mpi4py
# h5py
RUN pip install h5py
# other python stuff
RUN pip install wandb && \
pip install ruamel.yaml && \
pip install --upgrade tqdm && \
pip install timm && \
pip install einops
# benchy
RUN pip install git+https://github.com/romerojosh/benchy.git
# set wandb to offline
#ENV WANDB_MODE offline
# copy source code
RUN mkdir -p /opt/ERA5_wind
COPY config /opt/ERA5_wind/config
COPY copernicus /opt/ERA5_wind/copernicus
COPY docker /opt/ERA5_wind/docker
COPY networks /opt/ERA5_wind/networks
COPY utils /opt/ERA5_wind/utils
COPY plotting /opt/ERA5_wind/plotting
COPY mpu /opt/ERA5_wind/mpu
COPY *.py /opt/ERA5_wind/
COPY *.sh /opt/ERA5_wind/
COPY perf_tests /opt/perf_tests
# create dummy git image
RUN cd /opt/ERA5_wind && git init
docker/build.sh 0 → 100755
View file @ bdd87fae
#!/bin/bash
repo=gitlab-master.nvidia.com:5005/tkurth/era5_wind
#tag=latest
#tag=debug
tag=jaideep_legacy_dataloader
cd ../
# build
docker build -t ${repo}:${tag} -f docker/Dockerfile .
# push
docker push ${repo}:${tag}
# retag and repush
#docker tag ${repo}:${tag} thorstenkurth/era5-wind:${tag}
#docker push thorstenkurth/era5-wind:${tag}
export_DDP_vars.sh 0 → 100755
View file @ bdd87fae
export RANK=$SLURM_PROCID
export WORLD_RANK=$SLURM_PROCID
export LOCAL_RANK=$SLURM_LOCALID
export WORLD_SIZE=$SLURM_NTASKS
export MASTER_PORT=29500 # default from torch launcher
export WANDB_START_METHOD="thread"
inference/README_inference.md 0 → 100755
View file @ bdd87fae
## Inference scripts
To run inference from base dir:
* single member inference:
* inference.py for backbone variables: (see help for other options)
```bash
python inference/inference.py --config afno_backbone --run_num 0
```
* inference_precip.py for total precipitation: (see help for other options)
```bash
python inference/inference.py --config precip --run_num 0
```
* ensemble inference:
* inference_ensemble.py for wind variables: (see help for other options), use submit_batch_ensemble.sh to submit parallelized ensembles across
different initial conditions
* inference_ensemble_precip.py for total precipitation: (see help for other options), use submit_batch_ensemble.sh to submit parallelized ensembles across
different initial conditions; change to inference_ensemble_precip.py in the launch cmd
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment