schihei/pangu-weather
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fixed bugs in inference code

Browse Source
This commit is contained in:
Lingxi Xie
2023-03-02 22:37:27 +08:00
parent 0387b09086
commit cf3547114d
3 changed files with 4 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
inference_cpu.py
+1 -1
View File
@@ -21,7 +21,7 @@ options.intra_op_num_threads = 1
cuda_provider_options = {'arena_extend_strategy':'kSameAsRequested',} cuda_provider_options = {'arena_extend_strategy':'kSameAsRequested',}
# Initialize onnxruntime session for Pangu-Weather Models # 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 # Load the upper-air numpy arrays
input = np.load(os.path.join(input_data_dir, 'input_upper.npy')).astype(np.float32) input = np.load(os.path.join(input_data_dir, 'input_upper.npy')).astype(np.float32)
inference_gpu.py
+1 -1
View File
@@ -21,7 +21,7 @@ options.intra_op_num_threads = 1
cuda_provider_options = {'arena_extend_strategy':'kSameAsRequested',} cuda_provider_options = {'arena_extend_strategy':'kSameAsRequested',}
# Initialize onnxruntime session for Pangu-Weather Models # 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 # Load the upper-air numpy arrays
input = np.load(os.path.join(input_data_dir, 'input_upper.npy')).astype(np.float32) input = np.load(os.path.join(input_data_dir, 'input_upper.npy')).astype(np.float32)
inference_iterative.py
+2 -2
View File
@@ -22,8 +22,8 @@ options.intra_op_num_threads = 1
cuda_provider_options = {'arena_extend_strategy':'kSameAsRequested',} cuda_provider_options = {'arena_extend_strategy':'kSameAsRequested',}
# Initialize onnxruntime session for Pangu-Weather Models # 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)])
ort_session_6 = ort.InferenceSession('pangu_weather_6.onnx', sess_options=options, provider=[('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 # Load the upper-air numpy arrays
input = np.load(os.path.join(input_data_dir, 'input_upper.npy')).astype(np.float32) input = np.load(os.path.join(input_data_dir, 'input_upper.npy')).astype(np.float32)