Update pseudocode.py

README.md

@@ -18,10 +18,10 @@ The downloaded files shall be organized as the following hierarchy:
 │   │   ├── input_surface.npy
 │   │   ├── input_upper.npy
 │   ├── output_data
-│   ├── model_jit_cpu_1.onnx
-│   ├── model_jit_cpu_3.onnx
-│   ├── model_jit_cpu_6.onnx
-│   ├── model_jit_cpu_24.onnx
+│   ├── pangu_weather_1.onnx
+│   ├── pangu_weather_3.onnx
+│   ├── pangu_weather_6.onnx
+│   ├── pangu_weather_24.onnx
 │   ├── inference_cpu.py
 │   ├── inference_gpu.py
 │   ├── inference_iterative.py
@@ -41,15 +41,15 @@ pip install -r requirement_gpu.txt
 
 #### Downloading trained models
 
-Please download the four pre-trained models (~1.1GB each) from the Google drive:
+Please download the four pre-trained models (~1.1GB each) from Google drive or Baidu netdisk:
 
-The 1-hour model: [model_jit_cpu_1.onnx](https://drive.google.com/file/d/1fg5jkiN_5dHzKb-5H9Aw4MOmfILmeY-S/view?usp=share_link)
+The 1-hour model (pangu_weather_1.onnx): [Google drive](https://drive.google.com/file/d/1fg5jkiN_5dHzKb-5H9Aw4MOmfILmeY-S/view?usp=share_link)/[Baidu netdisk](https://pan.baidu.com/s/1M7SAigVsCSH8hpw6DE8TDQ?pwd=ie0h)
 
-The 3-hour model: [model_jit_cpu_3.onnx](https://drive.google.com/file/d/1EdoLlAXqE9iZLt9Ej9i-JW9LTJ9Jtewt/view?usp=share_link)
+The 3-hour model (pangu_weather_3.onnx): [Google drive](https://drive.google.com/file/d/1EdoLlAXqE9iZLt9Ej9i-JW9LTJ9Jtewt/view?usp=share_link)/[Baidu netdisk](https://pan.baidu.com/s/197fZsoiCqZYzKwM7tyRrfg?pwd=gmcl)
 
-The 6-hour model: [model_jit_cpu_6.onnx](https://drive.google.com/file/d/1a4XTktkZa5GCtjQxDJb_fNaqTAUiEJu4/view?usp=share_link)
+The 6-hour model (pangu_weather_6.onnx): [Google drive](https://drive.google.com/file/d/1a4XTktkZa5GCtjQxDJb_fNaqTAUiEJu4/view?usp=share_link)/[Baidu netdisk](https://pan.baidu.com/s/1q7IB7tNjqIwoGC7KVMPn4w?pwd=vxq3)
 
-The 24-hour model: [model_jit_cpu_24.onnx](https://drive.google.com/file/d/1lweQlxcn9fG0zKNW8ne1Khr9ehRTI6HP/view?usp=share_link)
+The 24-hour model (pangu_weather_24.onnx): [Google drive](https://drive.google.com/file/d/1lweQlxcn9fG0zKNW8ne1Khr9ehRTI6HP/view?usp=share_link)/[Baidu netdisk](https://pan.baidu.com/s/179q2gkz2BrsOR6g3yfTVQg?pwd=eajy)
 
 These models are stored using the ONNX format, and thus can be used via different languages such as Python, C++, C#, Java, etc.
 
@@ -69,11 +69,11 @@ We support ERA5 initial fields and ECMWF initial fields (e.g., the initial field
 
 We temporarily do not support other kinds of initial fields due to the possibly dramatic differences in the fields when Z<0.
 
-We provide an example of transferred input files, `input_surface.npy` and `input_upper.npy`, which correspond to the ERA5 initial fields of at 12:00UTC, 2018/09/27. Please download them using Google drive:
+We provide an example of transferred input files, `input_surface.npy` and `input_upper.npy`, which correspond to the ERA5 initial fields of at 12:00UTC, 2018/09/27. Please download them from Google drive or Baidu netdisk:
 
-[`input_surface.npy`](https://drive.google.com/file/d/1pj8QEVNpC1FyJfUabDpV4oU3NpSe0BkD/view?usp=share_link)
+`input_surface.npy`: [Google drive](https://drive.google.com/file/d/1pj8QEVNpC1FyJfUabDpV4oU3NpSe0BkD/view?usp=share_link)/[Baidu netdisk](https://pan.baidu.com/s/1i4o5i8guAqmOus6PWncAlA?pwd=4z9s)
 
-[`input_upper.npy`](https://drive.google.com/file/d/1--7xEBJt79E3oixizr8oFmK_haDE77SS/view?usp=share_link)
+`input_upper.npy`: [Google drive](https://drive.google.com/file/d/1--7xEBJt79E3oixizr8oFmK_haDE77SS/view?usp=share_link)/[Baidu netdisk](https://pan.baidu.com/s/1mS8X5MqEdbVfF2u2Us62FQ?pwd=sgx6)
 
 #### Inference
 
@@ -94,7 +94,11 @@ Note that one needs to download about 60TB of ERA5 data and prepare for computat
 
 ## License
 
-Pangu-Weather is released by Huawei Cloud. **The commercial use of these models is forbidden.**
+Pangu-Weather is released by Huawei Cloud.
+
+The trained parameters of Pangu-Weather are made available under the terms of the BY-NC-SA 4.0 license. You can find details [here](https://creativecommons.org/licenses/by-nc-sa/4.0/).
+
+**The commercial use of these models is forbidden.**
 
 Also, please note that all models were trained using the ERA5 dataset provided by ECMWF. Please do follow [their policy](https://apps.ecmwf.int/datasets/licences/copernicus/).
 

inference_cpu.py

@@ -21,7 +21,7 @@ options.intra_op_num_threads = 1
 cuda_provider_options = {'arena_extend_strategy':'kSameAsRequested',}
 
 # Initialize onnxruntime session for Pangu-Weather Models
-ort_session_24 = ort.InferenceSession('pangu_weather_24.onnx', sess_options=options, provider=['CPUExecutionProvider'])
+ort_session_24 = ort.InferenceSession('pangu_weather_24.onnx', sess_options=options, providers=['CPUExecutionProvider'])
 
 # Load the upper-air numpy arrays
 input = np.load(os.path.join(input_data_dir, 'input_upper.npy')).astype(np.float32)

inference_gpu.py

@@ -21,7 +21,7 @@ options.intra_op_num_threads = 1
 cuda_provider_options = {'arena_extend_strategy':'kSameAsRequested',}
 
 # Initialize onnxruntime session for Pangu-Weather Models
-ort_session_24 = ort.InferenceSession('pangu_weather_24.onnx', sess_options=options, provider=[('CUDAExecutionProvider', cuda_provider_options)])
+ort_session_24 = ort.InferenceSession('pangu_weather_24.onnx', sess_options=options, providers=[('CUDAExecutionProvider', cuda_provider_options)])
 
 # Load the upper-air numpy arrays
 input = np.load(os.path.join(input_data_dir, 'input_upper.npy')).astype(np.float32)

inference_iterative.py

@@ -22,8 +22,8 @@ options.intra_op_num_threads = 1
 cuda_provider_options = {'arena_extend_strategy':'kSameAsRequested',}
 
 # Initialize onnxruntime session for Pangu-Weather Models
-ort_session_24 = ort.InferenceSession('pangu_weather_24.onnx', sess_options=options, provider=[('CUDAExecutionProvider', cuda_provider_options)])
-ort_session_6 = ort.InferenceSession('pangu_weather_6.onnx', sess_options=options, provider=[('CUDAExecutionProvider', cuda_provider_options)])
+ort_session_24 = ort.InferenceSession('pangu_weather_24.onnx', sess_options=options, providers=[('CUDAExecutionProvider', cuda_provider_options)])
+ort_session_6 = ort.InferenceSession('pangu_weather_6.onnx', sess_options=options, providers=[('CUDAExecutionProvider', cuda_provider_options)])
 
 # Load the upper-air numpy arrays
 input = np.load(os.path.join(input_data_dir, 'input_upper.npy')).astype(np.float32)

pseudocode.py

@@ -13,6 +13,7 @@ from Your_AI_Library import Linear, Conv3d, Conv2d, ConvTranspose3d, ConvTranspo
 # GeLU: the GeLU activation function, see Pytorch API or Tensorflow API
 # DropOut: the dropout function, available in all deep learning libraries
 # DropPath: the DropPath function, see the implementation of vision-transformer, see timm pakage of Pytorch
+# A possible implementation of DropPath: from timm.models.layers import DropPath
 # LayerNorm: the layer normalization function, see Pytorch API or Tensorflow API
 # Softmax: softmax function, see Pytorch API or Tensorflow API
 from Your_AI_Library import GeLU, DropOut, DropPath, LayerNorm, SoftMax