search

naseemap47 / YOLO-NAS

Train and Inference your custom YOLO-NAS model by Single Command Line
Apache License 2.0
98 stars 13 forks source link

Training on two GPUs #21

Closed HaoqianSong closed 1 year ago

HaoqianSong commented 1 year ago

How to train a model using two GPUs simultaneously to improve speed, the current program only supports training on one GPU. How to modify the train.py file to enable training on two GPUs simultaneously?

naseemap47 commented 1 year ago

Hi @yangguanghaiann , I will look into it. I will update the code for you. Once I update, I will inform you.

naseemap47 commented 1 year ago

Hi @yangguanghaiann , I updated the training code, now it can support more than one GPUs. Only need to give --gpus args. It will automatically take all available GPUs. I think this will solve your issue. If have any other question and problem, Please let me know. Thank you.