YOLOv5 - Use GPU for inference

When the weight file training is configured, we open the reasoning file detect.py and find the following part:

insert image description hereThe statement in the red box in the figure is to set whether to use the GPU during detection.

1. Use GPU detection, that is, the default setting is 0 or 1, 2 (depending on the number of GPU blocks in the computer, it is not ruled out that some servers have multiple GPUs installed)

parser.add_argument('--device', default='0', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')

ps: If you set the GPU and report an error, it may be that the pytorch version is installed incorrectly. It is installed as the CPU version or is not compatible with cuda. ​​You need to reinstall pytorch.

2. Use CPU detection, that is, default is set to CPU

parser.add_argument('--device', default='CPU', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')


Guess you like

Origin blog.csdn.net/weixin_39357271/article/details/123271108