wenet-e2e / wenet

Production First and Production Ready End-to-End Speech Recognition Toolkit
https://wenet-e2e.github.io/wenet/
Apache License 2.0
3.87k stars 1.03k forks source link

[device args support] Replace `gpu` args with `device` args in train.py/recognize.py, etc. #2513

Open MengqingCao opened 1 month ago

MengqingCao commented 1 month ago

Is your feature request related to a problem? Please describe. I'm running wenet on Ascend NPU device recently. And I find that the scripts, e.g., train.py, recognize.py, ./cli/modle.py, only support GPU device via --gpu argument.

Describe the solution you'd like Specifying the accelerator to be used through --device argument will make it more convenient for people using wenet with different devices. There are a lot of devices (gpu, mps, xpu, etc.) supported in PyTorch, including Ascend NPU. If the replacement brings too many changes, adding a --device argumentcould also solve the problem, although I think the way of replacing is more reasonable.

Please let me know if there is any suggestions or concerns, thx!