新增調用 GPU 推理功能

[bensonbs]

調用 GPU 推理功能

測試以下作法能夠成功調用GPU推理

使用 onnxruntime gpu版本

pip install onnxruntime-gpu

測試是否安裝成功

import onnxruntime
print(onnxruntime.get_device())
GPU

測試是否調用成功

ort_session = onnxruntime.InferenceSession("./G2PWModel-v2-onnx/g2pw.onnx",providers=['CUDAExecutionProvider'])
print(ort_session.get_providers())
['CUDAExecutionProvider', 'CPUExecutionProvider']

理論上只須修改此行

g2pW/g2pw/api.py

Line 70 in 36b02d4

self.session_g2pw = onnxruntime.InferenceSession(os.path.join(model_dir, 'g2pw.onnx'), sess_options=sess_options)

self.session_g2pw =  onnxruntime.InferenceSession(os.path.join(model_dir, 'g2pw.onnx'), sess_options=sess_options,providers=['CUDAExecutionProvider'])